Schrödinger Equation Solver: Harmonic Oscillator

ShrödingerEquationSolver.swift

//: # Harmonic Oscillator Schrödinger Equation
//:
//: Numerically Solve for Shrödinger's Equation using the
//: Numerov algorithm.  This example inspired by Quantum
//: Physics for Dummies, Steven Holzner, John Wiley & Sons, Inc. 2013.
//:
//: Had to increase the delta (step size) and reduce the maximum allowed 
//: error to make it run in a reasonable amount of time in the playground.
//: (This obviously could be solved by switching to [more] compiled code.)
//:
//: The ideal result is a ground state of a proton is 1.51 MeV
//:
//: This is a very object oriented approach and could be made much more functional.
//: For instance, the psi values are stored in an array but technically don't need to be
//: without any loss in precision.

import Foundation

final class ShrödingerEquation {

  let Emin = 1.490
  let xMin = -15.0  // fm (fempto-meters)
  let xMax = 15.0   // fm
  let hZero: Double
  let EDelta = 0.0001
  let maxPsi = 0.00001
  let numberDivisions = 200
  
  var psi: [Double]
  var ECurrent: Double
  
  init() {
    ECurrent = Emin
    psi = Array(count: numberDivisions+1, repeatedValue: Double.NaN)
    psi[0] = 0
    psi[1] = -0.000000001
    psi[numberDivisions] = 1
    hZero = (xMax - xMin) / Double(numberDivisions)
  }

  func calculate() -> Double {
    
    while(fabs(psi[numberDivisions]) > maxPsi) {
      for i in 1..<numberDivisions {
        psi[i+1] = calculateNextPsi(i)
      }
      if psi[numberDivisions] > 0 {
        ECurrent -= EDelta
      }
      else {
        ECurrent += EDelta
      }
      print("𝝍: \(psi[numberDivisions]) E: \(ECurrent)")
    }

    print("Ground state energy: \(ECurrent) MeV")
    
    return ECurrent
  }
  
  func calculateKSquared(n: Int) -> Double {
    let x = hZero * Double(n) + xMin
    return (((0.05) * ECurrent) -  ((x * x) * 5.63e-3))
  }
  
  func calculateNextPsi(n: Int) -> Double {
    
    let KSqNMinusOne = calculateKSquared(n-1)
    let KSqN         = calculateKSquared(n)
    let KSqNPlusOne  = calculateKSquared(n+1)
    
    var nextPsi = 2.0 * (1.0 - (5.0 * hZero * hZero * KSqN / 12.0)) * psi[n]
    nextPsi -= (1 + (hZero * hZero * KSqNMinusOne / 12)) * psi[n-1]
    nextPsi /= 1 + (hZero * hZero * KSqNPlusOne / 12)
    return nextPsi
  }
}

//: Uncomment the below to make the computation run!
var se = ShrödingerEquation()
//se.calculate()