data/accretion_disk.py

import numpy as np


class AccretionDisk:
    def __init__(self, rs, inner_radius_rs, outer_radius_rs, brightness_scale, temperature_profile='power_law'):
        self.rs = rs
        self.r_inner = inner_radius_rs * rs
        self.r_outer = outer_radius_rs * rs
        self.brightness_scale = brightness_scale
        self.temperature_profile = temperature_profile

    def brightness(self, r, phi):
        if r < self.r_inner or r > self.r_outer:
            return 0.0
        r_norm = (r - self.r_inner) / (self.r_outer - self.r_inner)
        if self.temperature_profile == 'power_law':
            temp_factor = (r / self.r_inner) ** (-0.75)
            emission = temp_factor * (1.0 - np.sqrt(self.r_inner / r))
        elif self.temperature_profile == 'gaussian':
            peak = 0.3
            emission = np.exp(-0.5 * ((r_norm - peak) / 0.2) ** 2)
        else:
            emission = 1.0 / (r / self.r_inner) ** 2
        doppler_factor = 1.0 + 0.3 * np.sin(phi)
        return float(max(0.0, self.brightness_scale * emission * doppler_factor))

    def total_luminosity(self, num_samples=1000):
        r_samples = np.linspace(self.r_inner, self.r_outer, num_samples)
        total = sum(self.brightness(r, 0.0) * 2 * np.pi * r for r in r_samples)
        return total * (self.r_outer - self.r_inner) / num_samples


class ThinDisk(AccretionDisk):
    def __init__(self, rs, inner_radius_rs=3.0, outer_radius_rs=20.0, brightness_scale=1.0):
        super().__init__(rs, inner_radius_rs, outer_radius_rs, brightness_scale, 'power_law')


class ThickDisk(AccretionDisk):
    def __init__(self, rs, inner_radius_rs=2.0, outer_radius_rs=15.0, brightness_scale=1.0):
        super().__init__(rs, inner_radius_rs, outer_radius_rs, brightness_scale, 'gaussian')
        self.half_height_rs = 0.5 * rs

    def brightness(self, r, phi):
        base = super().brightness(r, phi)
        return base * 1.5