lensing.swift

// Integrate timelike geodesics in the equatorial plane (θ = π/2) of a Schwarzschild spacetime
// -----------------------------------------------------------------------------
// Units:  c = G = 1.  Mass M is given in the same geometric units (Schwarzschild radius r_s = 2M).
// Coordinates: (t, r, φ) with affine parameter λ.
// State vector y = (t, r, φ, 𝑡̇, ṙ, φ̇) where dots denote d/dλ.
// -----------------------------------------------------------------------------
import Foundation

/// Complete phase–space state of a particle on the geodesic
struct State {
    var t:   Double // coordinate time t
    var r:   Double // Schwarzschild radial coordinate r
    var phi: Double // azimuthal angle φ
    // four‑velocity components (derivatives w.r.t. affine parameter λ)
    var pt:   Double // dt/dλ
    var pr:   Double // dr/dλ
    var pphi: Double // dφ/dλ

    init(t: Double, r:Double, phi:Double, pt:Double, pr:Double, pphi:Double) {
        self.t = t; self.r = r; self.phi = phi; self.pt = pt; self.pr = pr; self.pphi = pphi
    }

    init(r:Double, phi:Double, pr:Double, pphi:Double, mass:Double) {
        self.t = 0.0; self.r = r; self.phi = phi; self.pr = pr; self.pphi = pphi
        self.pt = solvePt(r: r, pr: pr, pphi: pphi, mass: mass)
    }
}

typealias Derivative = State

/// Einstein‑summation geodesic right‑hand side using Christoffel symbols Γ^μ_{αβ}
/// Restricted to the equatorial plane θ = π/2, leaving θ constant.
func geodesicRHS(state y: State, mass M: Double) -> Derivative {
    let r   = y.r
    let pt  = y.pt
    let pr  = y.pr
    let pφ  = y.pphi

    // --- Christoffel symbols Γ^μ_{αβ} that are non‑zero in our 2+1 slice ---
    let γ_t_tr      =  M / (r * (r - 2*M))         // = Γ^t_{tr} = Γ^t_{rt}
    let γ_r_tt      =  M * (r - 2*M) / pow(r,3)    // Γ^r_{tt}
    let γ_r_rr      = -M / (r * (r - 2*M))         // Γ^r_{rr}
    let γ_r_φφ      = -(r - 2*M)                   // Γ^r_{φφ}  (remember metric signature)
    let γ_φ_rφ      = 1.0 / r                      // Γ^φ_{rφ} = Γ^φ_{φr}

    // --- First‑order pieces (position derivatives) ---
    let dt_dλ   = pt
    let dr_dλ   = pr
    let dφ_dλ   = pφ

    // --- Second‑order geodesic equations (velocity derivatives) ---
    let dpt_dλ  = -2.0 * γ_t_tr * pt * pr
    let dpr_dλ  = -γ_r_tt  * pow(pt,2) - γ_r_rr  * pow(pr,2) - γ_r_φφ  * pow(pφ,2)
    let dpφ_dλ  = -2.0 * γ_φ_rφ * pr * pφ

    return Derivative(t: dt_dλ, r: dr_dλ, phi: dφ_dλ,
                      pt: dpt_dλ, pr: dpr_dλ, pphi: dpφ_dλ)
}

/// Classic RK4 single integration step
typealias RHSFunction = (State, Double) -> Derivative
func rk4Step(state y: State, step h: Double, mass M: Double, rhs: RHSFunction) -> State {
    let k1 = rhs(y, M)

    let y2 = State(t:   y.t   + 0.5*h*k1.t,
                   r:   y.r   + 0.5*h*k1.r,
                   phi: y.phi + 0.5*h*k1.phi,
                   pt:  y.pt  + 0.5*h*k1.pt,
                   pr:  y.pr  + 0.5*h*k1.pr,
                   pphi:y.pphi + 0.5*h*k1.pphi)
    let k2 = rhs(y2, M)

    let y3 = State(t:   y.t   + 0.5*h*k2.t,
                   r:   y.r   + 0.5*h*k2.r,
                   phi: y.phi + 0.5*h*k2.phi,
                   pt:  y.pt  + 0.5*h*k2.pt,
                   pr:  y.pr  + 0.5*h*k2.pr,
                   pphi:y.pphi + 0.5*h*k2.pphi)
    let k3 = rhs(y3, M)

    let y4 = State(t:   y.t   + h*k3.t,
                   r:   y.r   + h*k3.r,
                   phi: y.phi + h*k3.phi,
                   pt:  y.pt  + h*k3.pt,
                   pr:  y.pr  + h*k3.pr,
                   pphi:y.pphi + h*k3.pphi)
    let k4 = rhs(y4, M)

    func combine(_ yVal: Double, _ k1: Double, _ k2: Double, _ k3: Double, _ k4: Double) -> Double {
        return yVal + h/6.0 * (k1 + 2.0*k2 + 2.0*k3 + k4)
    }

    return State(
        t:   combine(y.t,   k1.t,   k2.t,   k3.t,   k4.t),
        r:   combine(y.r,   k1.r,   k2.r,   k3.r,   k4.r),
        phi: combine(y.phi, k1.phi, k2.phi, k3.phi, k4.phi),
        pt:  combine(y.pt,  k1.pt,  k2.pt,  k3.pt,  k4.pt),
        pr:  combine(y.pr,  k1.pr,  k2.pr,  k3.pr,  k4.pr),
        pphi:combine(y.pphi,k1.pphi,k2.pphi,k3.pphi,k4.pphi)
    )
}

// -----------------------------------------------------------------------------
// MARK: – Helper to normalise the 4‑velocity for null geodesics
// -----------------------------------------------------------------------------
/// Given (r, pr, pphi) and mass M, returns the required pt ensuring g_{μν} u^μ u^ν = 0
func solvePt(r: Double, pr: Double, pphi: Double, mass M: Double) -> Double {
    let f = 1.0 - 2.0*M/r
    // Metric components in equatorial plane
    // g_tt = -(1 - 2M/r) = -f
    // g_rr = 1/f
    // g_φφ = r^2
    // −f pt² + pr² / f + r² pφ² = 0   (null geodesic)
    let ptSquared = ((pr*pr)/f + r*r*pphi*pphi) / f
    return sqrt(ptSquared)
}

struct CLI {
    var mass: Double = 0.5            // Black‑hole mass M (in geometric units)
    var r0: Double = 14.14             // Initial radius (must be > 2M)
    var phi0: Double = -0.78            // Initial azimuth
    var pr0: Double = -0.69             // Initial radial velocity d r / dλ
    var pphi0: Double = 0.052          // Initial angular velocity dφ / dλ
    var step: Double = 0.01           // Integration step Δλ
    var steps: Int = 7000             // Number of steps to compute

    mutating func parse() {
        let args = CommandLine.arguments.dropFirst() // remove executable name
        var index = args.startIndex
        while index < args.endIndex {
            switch args[index] {
            case "--mass":   index = args.index(after: index); mass  = Double(args[index]) ?? mass
            case "--r0":     index = args.index(after: index); r0    = Double(args[index]) ?? r0
            case "--phi0":   index = args.index(after: index); phi0  = Double(args[index]) ?? phi0
            case "--pr0":    index = args.index(after: index); pr0   = Double(args[index]) ?? pr0
            case "--pphi0":  index = args.index(after: index); pphi0 = Double(args[index]) ?? pphi0
            case "--step":   index = args.index(after: index); step  = Double(args[index]) ?? step
            case "--steps":  index = args.index(after: index); steps = Int(args[index])     ?? steps
            default:
                print("Unknown option \(args[index]) (ignored)")
            }
            index = args.index(after: index)
        }
    }
}

struct Pic {
    var W: Int
    var H: Int
    var pix : [UInt8]
    var stride: Int

    init(w: Int, h:Int) {
        W = w
        H = h
        stride = (w * 3 + 3) & ~3
        pix = [UInt8](repeating: 255, count: stride * h)
    }

    mutating func put(x:Int, y:Int, r:UInt8, g:UInt8, b:UInt8) {
        let o = stride * y + x * 3
        pix[o] = b
        pix[o+1] = g
        pix[o+2] = r
    }
}

func writeBMP(_ pic: Pic, _ fname:String) {
    let width = pic.W
    let height = pic.H
    let bytesPerPixel = 3  // 24-bit BMP (RGB)
    let rowPadding = (4 - (width * bytesPerPixel) % 4) % 4
    let rowSize = width * bytesPerPixel + rowPadding
    let pixelArraySize = rowSize * height
    let fileHeaderSize = 14
    let dibHeaderSize = 40
    let fileSize = fileHeaderSize + dibHeaderSize + pixelArraySize
    var data = Data()

    data.append(contentsOf: [
        0x42, 0x4D, // Signature 'BM'
    ])

    data.append(contentsOf: withUnsafeBytes(of: UInt32(fileSize).littleEndian, Array.init)) // File size
    data.append(contentsOf: [0x00, 0x00, 0x00, 0x00]) // Reserved
    data.append(contentsOf: withUnsafeBytes(of: UInt32(fileHeaderSize + dibHeaderSize).littleEndian, Array.init)) // Pixel data offset
    data.append(contentsOf: withUnsafeBytes(of: UInt32(dibHeaderSize).littleEndian, Array.init)) // Header size
    data.append(contentsOf: withUnsafeBytes(of: Int32(width).littleEndian, Array.init)) // Image width
    data.append(contentsOf: withUnsafeBytes(of: Int32(height).littleEndian, Array.init)) // Image height
    data.append(contentsOf: withUnsafeBytes(of: UInt16(1).littleEndian, Array.init)) // Planes
    data.append(contentsOf: withUnsafeBytes(of: UInt16(24).littleEndian, Array.init)) // Bits per pixel
    data.append(contentsOf: withUnsafeBytes(of: UInt32(0).littleEndian, Array.init)) // Compression (0 = BI_RGB)
    data.append(contentsOf: withUnsafeBytes(of: UInt32(pixelArraySize).littleEndian, Array.init)) // Image size
    data.append(contentsOf: withUnsafeBytes(of: Int32(2835).littleEndian, Array.init)) // X pixels per meter (~72 DPI)
    data.append(contentsOf: withUnsafeBytes(of: Int32(2835).littleEndian, Array.init)) // Y pixels per meter
    data.append(contentsOf: withUnsafeBytes(of: UInt32(0).littleEndian, Array.init)) // Colors in color table
    data.append(contentsOf: withUnsafeBytes(of: UInt32(0).littleEndian, Array.init)) // Important colors

    data.append(contentsOf: pic.pix)

    let url = URL(fileURLWithPath: fname)
    do {
        try data.write(to: url)
        print("BMP file created at \(url.path)")
    } catch {
        print("Error writing file: \(error)")
    }
}

// start
var cli = CLI()
cli.parse()

let W = 600, H = 600 // output image size
var pic = Pic(w:W, h:H)
let Pi = 3.14159265
let M = cli.mass

// draw the black hole in the middle
pic.put(x:W/2, y:H/2, r:0, g:0, b:0)
for a in 0 ..< 360 {
    let ang = Double(a)*Pi/180
    let x = W/2 + Int(10 * 2 * M * cos(ang))
    let y = H/2 + Int(10 * 2 * M * sin(ang))
    pic.put(x:x, y:y, r:120, g:120,b:120)     // horizon
    let x1 = W/2 + Int(10 * 3 * M * cos(ang))
    let y1 = H/2 + Int(10 * 3 * M * sin(ang))
    pic.put(x:x1, y:y1, r:220, g:220,b:220)   // photon sphere
}
let tk = 3.0
let N = 3000 // number of light rays
let tstep = 2.0
var movie = [Pic](repeating: pic, count: 256)
for n in 0..<N {  // cast N rays from the source
    let x0 = 7.0, y0 = 20.0 // light source coords
    let na = 2 * Pi * Double(n) / Double(N)
    let dx = sin(na), dy = -cos(na)
    let x1 = x0+dx, y1 = y0+dy
    let r0 = sqrt(x0*x0 + y0*y0)
    let ph0 = atan(y0/x0)
    let r1 = sqrt(x1*x1 + y1*y1)
    let ph1 = atan(y1/x1)
    var lastT = 0.0
    var lastFrame = -1

    var state = State(r: r0, phi: ph0, pr: r1-r0, pphi: ph1-ph0, mass: M)
    var λ = 0.0
    for _ in 0..<cli.steps {  // track points of a geodesic
        // stop at horizon or if time is too late to draw
        if state.r - 2 * cli.mass < 0.02 || state.t * tk > 255 {
            break
        }
        let x = W/2 + Int(state.r * cos(state.phi) * 10)
        let y = H/2 + Int(state.r * sin(state.phi) * 10)
        if x >= 0 && y >= 0 && x < W && y < H {
            if (n % 10)==0 || state.t - lastT >= tstep { // draw the grid in single pic
                let gg = UInt8( 128 + 120 * sin(state.t * Pi / 10) )
                let rr = UInt8(state.t * tk)
                pic.put(x:x, y:y, r:rr, g:gg, b:255)
                lastT = state.t
            }
            let frame = Int(state.t * 4)
            if (frame > lastFrame && frame < movie.count) { // draw in movie frames
                movie[frame].put(x:x, y:y, r:0, g:0, b:255)
                lastFrame = frame
            }
        }

        // Advance one RK4 step
        state = rk4Step(state: state, step: cli.step, mass: cli.mass, rhs: geodesicRHS)
        λ += cli.step
    }
}
writeBMP(pic, "geo.bmp")
for i in 0..<movie.count {
    writeBMP(movie[i], String(format: "frm%03d.bmp", i))
}