!pip install -I "phoebe>=2.0,<2.1"

%matplotlib inline

"""
Explanation: Rossiter-McLaughlin Effect
Setup
Let's first make sure we have the latest version of PHOEBE 2.0 installed. (You can comment out this line if you don't use pip for your installation or don't want to update to the latest release).
End of explanation
"""


import phoebe
from phoebe import u # units
import numpy as np
import matplotlib.pyplot as plt

logger = phoebe.logger()

b = phoebe.default_binary()

"""
Explanation: As always, let's do imports and initialize a logger and a new bundle.  See Building a System for more details.
End of explanation
"""


b['q'] = 0.7

b['rpole@primary'] = 1.0

b['rpole@secondary'] = 0.5

b['teff@secondary@component'] = 5000

"""
Explanation: Let's make a significant mass ratio and radius ratio...
End of explanation
"""


b['syncpar@primary@component'] = 2

"""
Explanation: Make sure the primary star is spinning quickly...
End of explanation
"""


b.add_dataset('rv', times=np.linspace(0,2,201), dataset='dynamicalrvs')  
# TODO: can't set rv_method here because compute options don't exist yet... and that's kind of annoying

b.add_dataset('rv', times=np.linspace(0,2,201), dataset='numericalrvs')

"""
Explanation: Adding Datasets
Now we'll add radial velocity and mesh datasets.  We'll add two identical datasets for RVs so that we can have one computed dynamically and the other computed numerically (these options will need to be set later).
End of explanation
"""


times = b.get_value('times@primary@numericalrvs@dataset')
times = times[times<0.1]
print times

b.add_dataset('mesh', dataset='mesh01', times=times)

"""
Explanation: Storing the mesh at every timestep is overkill, and will be both computationally and memory intensive.  So let's just sample at the times we care about.
End of explanation
"""


b.set_value_all('rv_method@dynamicalrvs@compute', 'dynamical')
b.set_value_all('rv_method@numericalrvs@compute', 'flux-weighted')

"""
Explanation: Running Compute
Now let's set the rv_method so that one dataset uses the dynamical method and the other uses the flux-weighted (numerical) method.  Note that here we have to use set_value_all or loop over the components, as technically there are parameters for each component-dataset pair.
End of explanation
"""


print b['rv_method']

b.run_compute(irrad_method='none')

"""
Explanation: Let's check to make sure that rv_method is set as we'd expect.
End of explanation
"""


axs, artists = b['dynamicalrvs@model'].plot(component='primary', color='b')
axs, artists = b['dynamicalrvs@model'].plot(component='secondary', color='r')

"""
Explanation: Plotting
Now let's plot the radial velocities.
First we'll plot the dynamical RVs.  Note that dynamical RVs show the true radial velocity of the center of mass of each star, and so we do not see the Rossiter McLaughlin effect.
End of explanation
"""


axs, artists = b['numericalrvs@model'].plot(component='primary', color='b')
axs, artists = b['numericalrvs@model'].plot(component='secondary', color='r')

"""
Explanation: But the numerical method integrates over the visible surface elements, giving us what we'd observe if deriving RVs from observed spectra of the binary.  Here we do see the Rossiter McLaughlin effect.  You'll also notice that RVs are not available for the secondary star when its completely occulted (they're nans in the array).
End of explanation
"""


fig = plt.figure(figsize=(12,12))
axs, artists = b['mesh@model'].plot(time=0.03, facecolor='rvs@numericalrvs', edgecolor=None)

"""
Explanation: To visualize what is happening, we can plot the radial velocities of each surface element in the mesh at one of these times.
Here just plot on the mesh@model parameterset - the mesh will automatically get coordinates from mesh01 and then we point to rvs@numericalrvs for the facecolors.
End of explanation
"""


fig = plt.figure(figsize=(12,12))
axs, artists = b['mesh01@model'].plot(time=0.09, facecolor='vzs', edgecolor=None)

"""
Explanation: Here you can see that the secondary star is blocking part of the "red" RVs of the primary star.
This is essentially the same as plotting the negative z-component of the velocities (for convention - our system is in a right handed system with +z towards the viewer, but RV convention has negative RVs for blue shifts).
We could also plot the RV per triangle by plotting 'vzs'.  Note that this is actually defaulting to an inverted colormap to show you the same colorscheme ('RdBu_r' vs 'RdBu').
End of explanation
"""


