Point Spread Function Modeling

class KeplerPRF:
    """
    Kepler Point Response Function model for PSF photometry.
    Provides accurate PSF shape modeling across the Kepler field.
    """
    
    def __init__(self, channel, shape, column, row):
        """
        Initialize Kepler PRF model.
        
        Parameters:
        - channel: int - Kepler CCD channel (1-84)
        - shape: tuple - Shape of PRF model (rows, cols)
        - column: float - Column position on detector
        - row: float - Row position on detector
        """
    
    def evaluate(self, column, row, flux, **kwargs):
        """
        Evaluate PRF model at given positions.
        
        Parameters:
        - column: float or array - Column positions
        - row: float or array - Row positions  
        - flux: float or array - Source fluxes
        
        Returns:
        2D array with evaluated PRF
        """

class SimpleKeplerPRF(KeplerPRF):
    """Simplified Kepler PRF model with reduced complexity"""

class GaussianPrior:
    """Gaussian prior distribution for Bayesian PRF fitting"""
    
    def __init__(self, mean, var):
        """
        Parameters:
        - mean: float - Prior mean
        - var: float - Prior variance
        """

class UniformPrior:
    """Uniform prior distribution"""
    
    def __init__(self, lower, upper):
        """
        Parameters:
        - lower: float - Lower bound
        - upper: float - Upper bound
        """

class FixedValuePrior:
    """Fixed value prior for known parameters"""
    
    def __init__(self, value):
        """
        Parameters:
        - value: float - Fixed parameter value
        """

class StarPrior:
    """Prior for stellar position and flux parameters"""

class BackgroundPrior:
    """Prior for background level parameters"""

class FocusPrior:
    """Prior for telescope focus parameters"""

class MotionPrior:
    """Prior for spacecraft motion parameters"""

class StarParameters:
    """Container for star-related model parameters"""
    
    def __init__(self, stars):
        """
        Parameters:
        - stars: list - List of star parameter dictionaries
        """

class BackgroundParameters:
    """Container for background model parameters"""

class FocusParameters:
    """Container for focus model parameters"""

class MotionParameters:
    """Container for motion model parameters"""

class TPFModelParameters:
    """Container for all target pixel file model parameters"""
    
    def __init__(self, stars=None, background=None, focus=None, motion=None):
        """
        Parameters:
        - stars: StarParameters - Star model parameters
        - background: BackgroundParameters - Background parameters
        - focus: FocusParameters - Focus parameters  
        - motion: MotionParameters - Motion parameters
        """

class TPFModel:
    """
    Complete target pixel file model for PRF photometry.
    Combines PRF, background, focus, and motion models.
    """
    
    def __init__(self, params):
        """
        Parameters:
        - params: TPFModelParameters - Model parameters
        """
    
    def evaluate(self, **kwargs):
        """Evaluate complete TPF model"""
    
    def gradient(self, **kwargs):
        """Compute model gradients for optimization"""

class PRFPhotometry:
    """
    High-level interface for PRF photometry analysis.
    Provides complete workflow from model setup to light curve extraction.
    """
    
    def __init__(self, tpf):
        """
        Initialize PRF photometry.
        
        Parameters:
        - tpf: TargetPixelFile - Target pixel file for analysis
        """
    
    def setup_model(self, star_priors=None, **kwargs):
        """
        Set up PRF model with priors.
        
        Parameters:
        - star_priors: list - Prior distributions for stars
        - background_prior: Prior - Background level prior
        - focus_prior: Prior - Focus parameter prior
        - motion_prior: Prior - Motion parameter prior
        
        Returns:
        Configured TPFModel
        """
    
    def fit(self, model=None, **kwargs):
        """
        Fit PRF model to data.
        
        Parameters:
        - model: TPFModel - Model to fit (uses setup_model if None)
        - method: str - Optimization method ('gradient', 'mcmc')
        - max_iterations: int - Maximum optimization iterations
        
        Returns:
        Fitted model parameters and uncertainties
        """
    
    def extract_lightcurve(self, fitted_params, **kwargs):
        """
        Extract light curve from fitted PRF model.
        
        Parameters:
        - fitted_params: Fitted model parameters
        - star_id: int - Which star to extract (for multi-star models)
        
        Returns:
        LightCurve object with PRF photometry
        """

import lightkurve as lk

# Download target pixel file
search = lk.search_targetpixelfile('Kepler-10')
tpf = search[0].download()

# Initialize PRF model
channel = tpf.channel
prf = lk.KeplerPRF(channel=channel, 
                   shape=tpf.flux[0].shape,
                   column=tpf.column,
                   row=tpf.row)

# Evaluate PRF at target position
prf_model = prf.evaluate(column=tpf.column + 2, 
                        row=tpf.row + 2,
                        flux=1000)

# Visualize PRF
import matplotlib.pyplot as plt
plt.imshow(prf_model, origin='lower')
plt.title('Kepler PRF Model')

# Note: Requires oktopus package
try:
    # Initialize PRF photometry
    prf_phot = lk.PRFPhotometry(tpf)
    
    # Set up priors for target star
    star_priors = [
        {
            'column': lk.GaussianPrior(tpf.column + 2, 0.1),
            'row': lk.GaussianPrior(tpf.row + 2, 0.1),
            'flux': lk.UniformPrior(500, 2000)
        }
    ]
    
    # Configure model
    model = prf_phot.setup_model(star_priors=star_priors)
    
    # Fit model to data
    fitted_params = prf_phot.fit(model, method='gradient')
    
    # Extract precise light curve
    prf_lc = prf_phot.extract_lightcurve(fitted_params)
    
    # Compare with simple aperture photometry
    aperture_lc = tpf.to_lightcurve()
    
    aperture_lc.plot(label='Aperture Photometry')
    prf_lc.plot(label='PRF Photometry')
    
    print(f"Aperture CDPP: {aperture_lc.estimate_cdpp():.1f} ppm")
    print(f"PRF CDPP: {prf_lc.estimate_cdpp():.1f} ppm")
    
except ImportError:
    print("Advanced PRF features require the oktopus package")
    print("Install with: pip install oktopus")

# Model multiple stars in the same TPF
try:
    # Define priors for multiple stars
    multi_star_priors = [
        {  # Primary target
            'column': lk.GaussianPrior(tpf.column + 2, 0.1),
            'row': lk.GaussianPrior(tpf.row + 2, 0.1), 
            'flux': lk.UniformPrior(800, 1200)
        },
        {  # Nearby contaminating star
            'column': lk.GaussianPrior(tpf.column + 4, 0.2),
            'row': lk.GaussianPrior(tpf.row + 1, 0.2),
            'flux': lk.UniformPrior(100, 400)
        }
    ]
    
    # Fit multi-star model
    prf_phot = lk.PRFPhotometry(tpf)
    model = prf_phot.setup_model(star_priors=multi_star_priors)
    fitted_params = prf_phot.fit(model)
    
    # Extract light curves for each star
    primary_lc = prf_phot.extract_lightcurve(fitted_params, star_id=0)
    contaminator_lc = prf_phot.extract_lightcurve(fitted_params, star_id=1)
    
    primary_lc.plot(label='Primary Star')
    contaminator_lc.plot(label='Contaminating Star')
    
except ImportError:
    print("Multi-star PRF modeling requires oktopus package")

# Analyze crowded field with multiple sources
try:
    # Download TPF for crowded region
    crowded_search = lk.search_targetpixelfile('M67', radius=0.1)
    crowded_tpf = crowded_search[0].download()
    
    # Set up many-star model (example with 5 stars)
    star_positions = [
        (5, 5), (7, 6), (3, 8), (9, 4), (6, 9)
    ]
    
    crowded_priors = []
    for col, row in star_positions:
        crowded_priors.append({
            'column': lk.GaussianPrior(col, 0.5),
            'row': lk.GaussianPrior(row, 0.5),
            'flux': lk.UniformPrior(50, 500)
        })
    
    # Fit crowded field model
    crowded_prf = lk.PRFPhotometry(crowded_tpf)
    crowded_model = crowded_prf.setup_model(star_priors=crowded_priors)
    crowded_fit = crowded_prf.fit(crowded_model, max_iterations=1000)
    
    # Extract individual light curves
    for star_id in range(len(star_positions)):
        star_lc = crowded_prf.extract_lightcurve(crowded_fit, star_id=star_id)
        star_lc.plot(label=f'Star {star_id+1}')
        
except ImportError:
    print("Crowded field analysis requires oktopus package")

# Advanced model configuration with all parameter types
try:
    prf_phot = lk.PRFPhotometry(tpf)
    
    # Define comprehensive priors
    star_priors = [{
        'column': lk.GaussianPrior(tpf.column + 2, 0.1),
        'row': lk.GaussianPrior(tpf.row + 2, 0.1),
        'flux': lk.UniformPrior(500, 1500)
    }]
    
    background_prior = lk.GaussianPrior(10, 5)  # Background level
    focus_prior = lk.GaussianPrior(1.0, 0.1)    # Focus parameter
    motion_prior = lk.GaussianPrior(0.0, 0.05)  # Motion parameter
    
    # Set up complete model
    model = prf_phot.setup_model(
        star_priors=star_priors,
        background_prior=background_prior,
        focus_prior=focus_prior,
        motion_prior=motion_prior
    )
    
    # Fit with MCMC for uncertainty estimation
    fitted_params = prf_phot.fit(model, method='mcmc')
    
    # Extract light curve with uncertainties
    prf_lc = prf_phot.extract_lightcurve(fitted_params)
    
except ImportError:
    print("Advanced PRF modeling requires oktopus package")