double_slit.swift

import Foundation

struct Complex {
    var re: Double
    var im: Double

    static func +(lhs: Complex, rhs: Complex) -> Complex {
        Complex(re: lhs.re + rhs.re, im: lhs.im + rhs.im)
    }

    static func -(lhs: Complex, rhs: Complex) -> Complex {
        Complex(re: lhs.re - rhs.re, im: lhs.im - rhs.im)
    }

    static func *(lhs: Complex, rhs: Complex) -> Complex {
        Complex(
            re: lhs.re * rhs.re - lhs.im * rhs.im,
            im: lhs.re * rhs.im + lhs.im * rhs.re
        )
    }

    static func *(lhs: Double, rhs: Complex) -> Complex {
        Complex(re: lhs * rhs.re, im: lhs * rhs.im)
    }

    static func *(lhs: Complex, rhs: Double) -> Complex {
        Complex(re: lhs.re * rhs, im: lhs.im * rhs)
    }

    func modulusSquared() -> Double {
        re * re + im * im
    }

    func ampl() -> Double { sqrt(modulusSquared()) }
}

let N = 200            // Grid size
let dx = 1.0
let dt = 0.01

let V_wall = 200.0
let slit_width = 3
let slit_gap = 30
let absorbWidth = 10

let wallX = 60
let centerY = N / 2
let tubeY = 35

var ψ = Array(repeating: Array(repeating: Complex(re: 0, im: 0), count: N), count: N)
var V_real = Array(repeating: Array(repeating: 0.0, count: N), count: N)
var V_imag = Array(repeating: Array(repeating: 0.0, count: N), count: N)

for y in 0..<N {
    V_real[wallX][y] = V_wall
}
for x in wallX..<N {
    V_real[x][tubeY] = V_wall
    V_real[x][N-tubeY] = V_wall
}
for offset in -slit_width..<slit_width {
    V_real[wallX][centerY - slit_gap / 2 + offset] = 0.0
    V_real[wallX][centerY + slit_gap / 2 + offset] = 0.0
}

// Absorbing boundary (imaginary potential)
for x in 0..<N {
    for y in 0..<N {
        var absorb = 0.0
        if x < absorbWidth {
            absorb += pow(Double(absorbWidth - x) / Double(absorbWidth), 2)
        }
        if x >= N - absorbWidth {
            absorb += pow(Double(x - (N - absorbWidth)) / Double(absorbWidth), 2)
        }
        if y < absorbWidth {
            absorb += pow(Double(absorbWidth - y) / Double(absorbWidth), 2)
        }
        if y >= N - absorbWidth {
            absorb += pow(Double(y - (N - absorbWidth)) / Double(absorbWidth), 2)
        }
        V_imag[x][y] = 5.0 * absorb
    }
}

// Gaussian wave packet
let σ = 2.0
let x0 = 20.0
let y0 = Double(N) / 2.0
let kx = 1.0

for x in 0..<N {
    for y in 0..<N {
        let dx2 = pow(Double(x) - x0, 2)
        let dy2 = pow(Double(y) - y0, 2)
        let envelope = exp(-(dx2 + dy2) / (2 * σ * σ))
        let phase = Complex(re: cos(kx * Double(x)), im: sin(kx * Double(x)))
        ψ[x][y] = envelope * phase
    }
}

func laplacian(x: Int, y: Int, ψ: [[Complex]]) -> Complex {
    let center = ψ[x][y]
    let left   = x > 0     ? ψ[x - 1][y] : Complex(re: 0, im: 0)
    let right  = x < N - 1 ? ψ[x + 1][y] : Complex(re: 0, im: 0)
    let up     = y > 0     ? ψ[x][y - 1] : Complex(re: 0, im: 0)
    let down   = y < N - 1 ? ψ[x][y + 1] : Complex(re: 0, im: 0)
    return (left + right + up + down - 4 * center) * (1.0 / (dx * dx))
}

func schrodingerRHS(ψ: [[Complex]]) -> [[Complex]] {
    var dψ = Array(repeating: Array(repeating: Complex(re: 0, im: 0), count: N), count: N)
    for x in 1..<N-1 {
        for y in 1..<N-1 {
            let lap = laplacian(x: x, y: y, ψ: ψ)
            let Vx = V_real[x][y]
            let Vi = V_imag[x][y]
            let Hψ = Complex(re: 0, im: 0.5) * lap + Complex(re: 0, im: -Vx) * ψ[x][y]
            let damping = Complex(re: -Vi, im: 0) * ψ[x][y]
            dψ[x][y] = Hψ + damping
        }
    }
    return dψ
}

func addScaled(_ ψ: [[Complex]], _ k: [[Complex]], scale: Double) -> [[Complex]] {
    var result = ψ
    for x in 0..<N {
        for y in 0..<N {
            result[x][y] = ψ[x][y] + scale * k[x][y]
        }
    }
    return result
}

// RK4 Time Evolution
func rk4Step(_ ψ: [[Complex]]) -> [[Complex]] {
    let k1 = schrodingerRHS(ψ: ψ)
    let k2 = schrodingerRHS(ψ: addScaled(ψ, k1, scale: 0.5 * dt))
    let k3 = schrodingerRHS(ψ: addScaled(ψ, k2, scale: 0.5 * dt))
    let k4 = schrodingerRHS(ψ: addScaled(ψ, k3, scale: dt))

    var ψ_next = ψ
    for x in 0..<N {
        for y in 0..<N {
            let update = (1.0 / 6.0) * (k1[x][y] + 2 * k2[x][y] + 2 * k3[x][y] + k4[x][y])
            ψ_next[x][y] = ψ[x][y] + dt * update
        }
    }
    return ψ_next
}

func complexToRGB(_ z: Complex, amplitudeScale: Double = 1.0) -> (UInt8, UInt8, UInt8) {
    let amplitude = min(1.0, z.ampl() * amplitudeScale)
    var phase = atan2(z.im, z.re)
    if phase < 0 { phase += 2 * Double.pi }  // normalize to [0, 2π]

    let hue = phase / (2 * Double.pi)  // [0, 1]
    let saturation = 1.0
    let value = amplitude              // brightness

    // HSV to RGB conversion
    let h = hue * 6
    let i = Int(h)
    let f = h - Double(i)
    let p = value * (1 - saturation)
    let q = value * (1 - saturation * f)
    let t = value * (1 - saturation * (1 - f))

    var r = 0.0, g = 0.0, b = 0.0
    switch i % 6 {
    case 0: r = value; g = t; b = p
    case 1: r = q; g = value; b = p
    case 2: r = p; g = value; b = t
    case 3: r = p; g = q; b = value
    case 4: r = t; g = p; b = value
    case 5: r = value; g = p; b = q
    default: break
    }

    return (
        UInt8(min(255, max(0, r * 255))),
        UInt8(min(255, max(0, g * 255))),
        UInt8(min(255, max(0, b * 255)))
    )
}

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: 0, 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 wavefunctionToPic(_ ψ: [[Complex]]) -> Pic {
    let H = ψ.count
    let W = ψ[0].count
    var pic = Pic(w: W+50, h: H)

    var maxAmplitude = 0.0
    for row in ψ {
        for z in row {
            let amp = z.ampl()
            if amp > maxAmplitude {
                maxAmplitude = amp
            }
        }
    }
    let scale = maxAmplitude > 0 ? 1.0 / maxAmplitude : 1.0

    for x in absorbWidth...W-30 {
        if x == W-30 {
            let probs = (absorbWidth..<H-absorbWidth).map { y in ψ[x][y].modulusSquared() }
            let maxProb = probs.max() ?? 1.0
            
            for y in absorbWidth..<H-absorbWidth {
                pic.put(x: x, y: y, r: 64, g: 64, b: 164)
                let a = probs[y - absorbWidth] / maxProb
                let c = UInt8(min(255, max(0, a * 255)))
                for i in -29..<50 {
                    pic.put(x: W+i, y:y, r:c, g:c, b:c)
                }
            }
        } else {
            for y in absorbWidth..<H-absorbWidth {
                let (r, g, b) = complexToRGB(ψ[x][y], amplitudeScale: scale)
                pic.put(x: x, y: y, r: r, g: g, b: b)
                if V_real[x][y] > 0.0 { pic.put(x: x, y: y, r: 228, g: 228, b: 228) }
            }
        }
    }

    return pic
}

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)")
    }
}

let steps = 40000
var np = 0
for step in 0..<steps {
    if step % 50 == 0 {
        var totalProb = 0.0
        for x in 0..<N {
            for y in 0..<N {
                totalProb += ψ[x][y].modulusSquared()
            }
        }
        print("Step \(step): Total probability = \(totalProb)")
        let pic = wavefunctionToPic(ψ)
        writeBMP(pic, String(format: "wf%04d.bmp", np))
        np += 1
    }
    ψ = rk4Step(ψ)
}