Implementing a simple exoplanet transit model mean function  #94

import numpy as np 
import matplotlib.pyplot as plt
import batman 

import jax
import jax.numpy as jnp

import tinygp
import jaxopt

from functools import partial
from tinygp import kernels, GaussianProcess

jax.config.update("jax_enable_x64", True)

### Keys of all transit parameters
all_transit_params = ['per', 'ecc', 'w', 'limb_dark', 't0', 'a', 'inc', 'u','rp']

### Dictionary of fixed transit parameters 
fix_transit_params = {
               "per":3.52474859,
             "ecc":0.000,
             "w":90.,
             "limb_dark":'quadratic',
             "t0":0.,
         "a":8.807,
        "inc": 86.744, 
    "u": [0.41739267, 0.15299455],
           }

### Dictionary of free transit parameters 
free_transit_params = {
    "rp": 0.12,
    "mean_level": 0.
}

### GP hyperparameters
params_gp = {
    "log_gp_amp":np.log(0.1),
    "log_gp_scale":np.log(3.0),
    "log_gp_diag":np.log(0.001),

}

### Merge the free transit params and the GP params 
free_params_all = {**free_transit_params, **params_gp}

### Function to call batman and calculate exoplanet transit model 
def tmodel(params,time):
    params_bat = batman.TransitParams()
    params_bat.__dict__.update(params)
    flux_bat = batman.TransitModel(params_bat, time.val ).light_curve(params_bat)
    return flux_bat
    

### Function to calculate the GP mean function 
def mean_function(params,X):
    """
    params: Dictionary of all free parameters.
    X: (X_grid) here time
    """
    time = X
    mean_level = params["mean_level"]
    
    params_in_func = {}
    
    for tparam in all_transit_params:
        if tparam in fix_transit_params.keys():
            params_in_func[tparam] = fix_transit_params[tparam]
        else:
            params_in_func[tparam] = params[tparam]
    
    lc = tmodel(params_in_func,time) + mean_level
        
    return lc

transit_duration = 3./24. ### in days 
X_grid = np.linspace(-transit_duration*2., transit_duration*2., 100)
print("X_grid shape", X_grid.shape)

model = jax.vmap(partial(mean_function, free_params_all))(X_grid)
print("model shape", model.shape)