Grid and Stellar Functions

from dustpy.std import grid

def grid.OmegaK(sim):
    """
    Calculate Keplerian orbital frequency.

    Computes the Keplerian frequency at each radial grid point
    based on the stellar mass and radial position. This is
    fundamental for calculating orbital time scales and
    gas pressure support.

    Parameters:
    - sim: Simulation object

    Returns:
    numpy.ndarray: Keplerian frequency [1/s] at each radial point

    Formula:
    Ω_K = √(GM_star / r³)

    Where:
    - G: Gravitational constant
    - M_star: Stellar mass
    - r: Radial distance from star
    """

from dustpy.std import star

def star.luminosity(sim):
    """
    Calculate stellar luminosity.

    Computes the stellar luminosity based on the stellar
    mass, radius, and effective temperature using the
    Stefan-Boltzmann law.

    Parameters:
    - sim: Simulation object

    Returns:
    float: Stellar luminosity [erg/s]

    Formula:
    L = 4π R²_star σ_SB T⁴_eff

    Where:
    - R_star: Stellar radius
    - σ_SB: Stefan-Boltzmann constant
    - T_eff: Effective temperature
    """

from dustpy import Simulation
from dustpy.std import grid
import dustpy.constants as c

# Create simulation with custom grid
sim = Simulation()

# Set grid parameters
sim.ini.grid.Nr = 200                   # Number of radial points
sim.ini.grid.rmin = 0.1 * c.au          # Inner radius
sim.ini.grid.rmax = 1000 * c.au         # Outer radius

# Create grids
sim.makegrids()

# Calculate Keplerian frequency
omega_k = grid.OmegaK(sim)              # [1/s]

# Convert to useful units
orbital_periods = 2 * c.pi / omega_k    # [s]
periods_years = orbital_periods / c.year # [years]

print(f"Orbital period at 1 AU: {periods_years[50]:.1f} years")
print(f"Orbital period at 10 AU: {periods_years[150]:.1f} years")

from dustpy.std import star
import dustpy.constants as c

# Set stellar parameters
sim.ini.star.M = 0.5 * c.M_sun          # Half solar mass
sim.ini.star.R = 0.7 * c.R_sun          # 70% solar radius
sim.ini.star.T = 4500                   # Effective temperature [K]

# Initialize simulation
sim.makegrids()
sim.initialize()

# Calculate luminosity
luminosity = star.luminosity(sim)       # [erg/s]
solar_luminosity = 3.828e33             # [erg/s]
L_ratio = luminosity / solar_luminosity

print(f"Stellar luminosity: {luminosity:.2e} erg/s")
print(f"Luminosity ratio (L/L_sun): {L_ratio:.2f}")

# Access grid properties
radial_centers = sim.grid.r             # [cm]
radial_interfaces = sim.grid.ri         # [cm]
annulus_areas = sim.grid.A              # [cm²]

# Calculate orbital frequencies
omega_k = grid.OmegaK(sim)              # [1/s]

# Derived time scales
dynamical_time = 1.0 / omega_k          # [s]
hill_sphere_radius = radial_centers * (sim.star.M / (3 * sim.star.M))**(1/3)

# Print grid information
print(f"Grid extends from {sim.grid.r[0]/c.au:.2f} to {sim.grid.r[-1]/c.au:.1f} AU")
print(f"Number of grid points: {len(sim.grid.r)}")
print(f"Innermost dynamical time: {dynamical_time[0]/c.year:.2e} years")

# Time-dependent stellar properties
def evolve_stellar_properties(sim, time):
    """Update stellar properties as function of time."""

    # Example: Simple stellar evolution model
    age_myr = time / (1e6 * c.year)      # Age in Myr

    # Luminosity decreases with time for low-mass stars
    if age_myr < 10:
        L_factor = 1.5 * (1 + np.exp(-age_myr/2))
    else:
        L_factor = 1.0

    # Update stellar luminosity
    base_luminosity = 4 * c.pi * sim.star.R**2 * c.sigma_sb * sim.star.T**4
    sim.star.L = base_luminosity * L_factor

    return sim.star.L

# Apply stellar evolution
current_time = 5e6 * c.year             # 5 Myr
new_luminosity = evolve_stellar_properties(sim, current_time)
print(f"Evolved luminosity: {new_luminosity:.2e} erg/s")

def analyze_grid_resolution(sim):
    """Analyze grid resolution and recommend improvements."""

    r = sim.grid.r                      # [cm]
    omega_k = grid.OmegaK(sim)          # [1/s]

    # Grid spacing
    dr = np.diff(sim.grid.ri)           # [cm]
    relative_spacing = dr[:-1] / r[1:]   # Relative grid spacing

    # Hill sphere resolution
    hill_radius = r * (sim.star.M / (3 * sim.star.M))**(1/3)
    points_per_hill = hill_radius / dr

    # Report
    print(f"Average relative spacing: {np.mean(relative_spacing):.3f}")
    print(f"Maximum relative spacing: {np.max(relative_spacing):.3f}")
    print(f"Minimum Hill sphere resolution: {np.min(points_per_hill):.1f} points")

    # Recommendations
    if np.max(relative_spacing) > 0.1:
        print("Warning: Grid spacing may be too coarse")
    if np.min(points_per_hill) < 5:
        print("Warning: Hill sphere under-resolved")

# Check grid quality
analyze_grid_resolution(sim)

def calculate_epicyclic_frequency(sim):
    """Calculate epicyclic frequency for eccentric orbits."""

    # For Keplerian potential, epicyclic frequency equals Keplerian
    omega_k = grid.OmegaK(sim)
    kappa = omega_k  # Epicyclic frequency

    return kappa

def calculate_vertical_frequency(sim):
    """Calculate vertical oscillation frequency."""

    # For thin disk approximation
    omega_k = grid.OmegaK(sim)
    nu_z = omega_k  # Vertical frequency ≈ Keplerian frequency

    return nu_z

# Calculate additional frequencies
epicyclic_freq = calculate_epicyclic_frequency(sim)
vertical_freq = calculate_vertical_frequency(sim)

# Oscillation periods
epicyclic_period = 2 * c.pi / epicyclic_freq
vertical_period = 2 * c.pi / vertical_freq

# Combine grid and stellar functions with disk physics
from dustpy.std import gas

# Initialize simulation
sim.makegrids()
sim.initialize()

# Calculate grid-dependent quantities
omega_k = grid.OmegaK(sim)              # [1/s]
luminosity = star.luminosity(sim)       # [erg/s]

# Use in disk temperature calculation
temperature = (luminosity / (16 * c.pi * c.sigma_sb * sim.grid.r**2))**0.25
sim.gas.T[:] = temperature

# Calculate pressure scale height
cs = gas.cs_isothermal(sim)             # [cm/s]
scale_height = cs / omega_k             # [cm]

print(f"Temperature at 1 AU: {temperature[50]:.0f} K")
print(f"Scale height at 1 AU: {scale_height[50]/c.au:.3f} AU")

def grid_diagnostics(sim):
    """Comprehensive grid quality diagnostics."""

    r = sim.grid.r
    ri = sim.grid.ri

    print("Grid Diagnostics:")
    print(f"  Radial range: {r[0]/c.au:.2f} - {r[-1]/c.au:.0f} AU")
    print(f"  Number of points: {len(r)}")

    # Grid spacing
    dr = np.diff(ri)
    print(f"  Grid spacing: {dr[0]/c.au:.4f} - {dr[-1]/c.au:.2f} AU")

    # Logarithmic spacing check
    r_ratio = r[1:] / r[:-1]
    is_logarithmic = np.allclose(r_ratio, r_ratio[0], rtol=0.01)
    print(f"  Logarithmic spacing: {is_logarithmic}")

    if is_logarithmic:
        print(f"  Log spacing factor: {r_ratio[0]:.3f}")

    # Time scale resolution
    omega_k = grid.OmegaK(sim)
    min_period = 2 * c.pi / np.max(omega_k) / c.year
    max_period = 2 * c.pi / np.min(omega_k) / c.year

    print(f"  Orbital periods: {min_period:.1f} - {max_period:.0f} years")

# Run diagnostics
grid_diagnostics(sim)