LambertWFromJulia.kt

/**
 * Source: https://github.com/JuliaMath/LambertW.jl/blob/master/src/LambertW.jl
 */
object LambertW {

	// Lambert W function for real z
	fun lambertW(z: Double, k: Int = 0, maxIts: Int = 1000): Double {
		return _lambertW(z, k, maxIts)
	}

	// Private function for computing Lambert W
	private fun _lambertW(z: Double, k: Int, maxIts: Int): Double {
		require(k == 0 || k == -1) { "For real z, k must be 0 or -1" }
		if (z.isNaN()) return Double.NaN
		if (z == Double.POSITIVE_INFINITY) return Double.POSITIVE_INFINITY
		if (z == 0.0 && k == 0) return 0.0
		if (z == 0.0 && k == -1) return Double.NEGATIVE_INFINITY
		if (k == 0) return lambertWBranchZero(z, maxIts)
		if (k == -1) return lambertWBranchMinusOne(z, maxIts)
		throw IllegalArgumentException("Invalid k value")
	}

	// Lambert W branch for k = 0
	private fun lambertWBranchZero(x: Double, maxIts: Int): Double {
		if (x == 0.0) return -1.0
		val oneOverE = -1.0 / E
		if (x <= oneOverE) return -1.0
		if (x > 1.0) {
			val lx = ln(x)
			val llx = ln(lx)
			val x0 = lx - llx - ln(-llx / lx + 1) * (1.0 / 2.0)
			return lambertWRootFinding(x, x0, maxIts)
		}
		else {
			val x0 = 0.567 * x // Magic number from Julia code
			return lambertWRootFinding(x, x0, maxIts)
		}
	}

	// Lambert W branch for k = -1
	private fun lambertWBranchMinusOne(x: Double, maxIts: Int): Double {
		val oneOverE = -1.0 / E
		if (x <= oneOverE) return -1.0
		if (x == 0.0) return Double.NEGATIVE_INFINITY
		if (x < 0.0) {
			return lambertWRootFinding(x, ln(-x), maxIts)
		}
		else {
			throw IllegalArgumentException("For k = -1, x must be less than 0")
		}
	}

	// Root finding using Halley's method
	private fun lambertWRootFinding(z: Double, x0: Double, maxIts: Int): Double {
		var x = x0
		var lastX: Double = x
		var lastDiff = 0.0
		var converged = false
		for (i in 1..maxIts) {
			val ex = exp(x)
			val xExZ = x * ex - z
			val x1 = x + 1
			x -= xExZ / (ex * x1 - (x + 2) * xExZ / (2 * x1))
			val xDiff = abs(lastX - x)
			if (xDiff <= 3 * Math.ulp(abs(lastX)) || lastDiff == xDiff) {
				converged = true
				break
			}
			lastX = x
			lastDiff = xDiff
		}
		if (!converged) {
			println("Warning: Lambert W function did not converge in $maxIts iterations.")
		}
		return x
	}
}

LambertWFromPariGp.kt

import kotlin.math.ln
import kotlin.math.sqrt
import kotlin.math.E
import kotlin.math.ln1p

//  https://pari.math.u-bordeaux.fr/git/pari.git

fun main() {
	println(Double.MAX_VALUE)
	println(lambertWm1Approx(-1 / E))
}

/**
 * Rough approximation to W0(a > -1/e), < 1% relative error
 */
fun lambertW0Approx(x: Double): Double {
	if (x < -0.2583) {
		val c2 = -1.0 / 3
		val c3 = 11.0 / 72
		val c4 = -43.0 / 540
		val c5 = 769.0 / 17280
		val p = sqrt(2 * (E * x + 1))
		return if (x < -0.3243) {
			-1 + p * (1 + p * (c2 + p * c3))
		} else {
			-1 + p * (1 + p * (c2 + p * (c3 + p * (c4 + p * c5))))
		}
	} else {
		var wd = ln1p(x)
		wd *= (1 - ln(wd / x)) / (1 + wd)
		return if (x < 0.6482 && x > -0.1838) {
			wd
		} else {
			wd * (1 - ln(wd / x)) / (1 + wd)
		}
	}
}

/**
 * Rough approximation to W_{-1}(0 > a > -1/e), < 1% relative error
 */
fun lambertWm1Approx(x: Double): Double {
	if (x < -0.2464) {
		val c2 = -1.0 / 3
		val c3 = 11.0 / 72
		val c4 = -43.0 / 540
		val c5 = 769.0 / 17280
		val p = -sqrt(2 * (E * x + 1))
		return if (x < -0.3243) {
			-1 + p * (1 + p * (c2 + p * c3))
		} else {
			-1 + p * (1 + p * (c2 + p * (c3 + p * (c4 + p * c5))))
		}
	} else {
		var wd: Double
		val negativeX = -x
		wd = -ln(negativeX)
		wd *= (1 - ln(wd / negativeX)) / (1 - wd)
		return if (negativeX < 0.0056) {
			-wd
		} else {
			-wd * (1 - ln(wd / negativeX)) / (1 - wd)
		}
	}
}

NewtonRaphson.kt

import kotlin.math.*

fun main() {
	val input = -0.1
	val initialApproximation = 1.0

	println(lambertW0(-0.1))
	println(lambertWm1(-0.1))
	println()

	val result2 = newtonRaphson(
		initialApproximation = -2.3,
		function = { x ->
			x * exp(x) - input
		},
		derivative = { x ->
			(x + 1) * exp(x)
		},
	)
	println(result2)

	val result3 = newtonRaphson(initialApproximation) { x ->
		x * exp(x) - input
	}
	println(result3)
}

fun lambertWm1Approx(a: Double): Double {
	if (a < -0.2464) {
		val c2 = -1.0 / 3
		val c3 = 11.0 / 72
		val c4 = -43.0 / 540
		val c5 = 769.0 / 17280
		val p = -sqrt(2 * (E * a + 1))
		return if (a < -0.3243) {
			-1 + p * (1 + p * (c2 + p * c3))
		} else {
			-1 + p * (1 + p * (c2 + p * (c3 + p * (c4 + p * c5))))
		}
	} else {
		val aModified = -a
		var wd = -ln(aModified)
		wd *= (1 - ln(wd / aModified)) / (1 - wd)
		return if (aModified < 0.0056) {
			-wd
		} else {
			-wd * (1 - ln(wd / aModified)) / (1 - wd)
		}
	}
}

fun lambertW0(
	x: Double,
	maxIterations: Int = 1000,
): Double {
	if (x < -1 / E) {
		throw IllegalArgumentException("Input ($x) can't be less than -1 / E")
	}

	when (x) {
		0.0 -> return 0.0
		-1.0 / E -> return -1.0
		1.0 -> return 0.5671432904097838
		E -> return 1.0
		// e * e ^ e
		41.19355567471612 -> return E
		Double.POSITIVE_INFINITY -> Double.POSITIVE_INFINITY
	}

	// Obtained from experimenting
	val initialApproximation = if (x <= 30) {
		val ln1pOfX = ln1p(x)
		(ln(1 + ln1pOfX) + ln1pOfX) / 2
	}
	else {
		// https://mathoverflow.net/questions/57819/best-approximation-to-the-lambertwx-or-explambertwx
		val lnOfX = ln(x)
		val lnOfLnOfX = ln(lnOfX)
		lnOfX - lnOfLnOfX + lnOfLnOfX / lnOfX
	}

	var approx = initialApproximation

	for (i in 1..maxIterations) {
		val previousApprox = approx
		val expResult = exp(approx)
		approx -= (approx * expResult - x) / (expResult * (approx + 1))

		if (approx == previousApprox) {
			return approx
		}
	}
	return approx
}

fun lambertWm1(
	x: Double,
	maxIterations: Int = 1000,
): Double {
	if (x > 0.0) {
		throw IllegalArgumentException("Input ($x) can't be greater than 0")
	}
	if (x < -1 / E) {
		throw IllegalArgumentException("Input ($x) can't be less than -1 / E")
	}

	if (x == 0.0) {
		return Double.NEGATIVE_INFINITY
	}
	if (x == -1 / E) {
		return -1.0
	}
	// -ln(2) / 2
	if (x == -0.34657359027997264) {
		// -ln(4)
		return -1.3862943611198906
	}


	val initialApproximation: Double = when (x) {
		in -1 / E..-0.364 -> {
			-1 - (1.4142135623730951 * sqrt(1 + E * x))
		}
		in -0.364..-0.3536 -> {
			-14.9347 * x - 6.595
		}
		in -0.3536..-0.3318 -> {
			-10 * x - 4.85
		}
		in -0.3318..-0.2762 -> {
			-7.89616 * x - 4.152
		}
		in -0.2762..-0.1622 -> {
			-8 * x - 4.155
		}
		else -> ln(-x / (2 * PI))
	}

	var approx = initialApproximation

	for (i in 1..maxIterations) {
		val previousApprox = approx
		val expResult = exp(approx)
		approx -= (approx * expResult - x) / (expResult * (approx + 1))

		if (approx == previousApprox) {
			return approx
		}
	}
	return approx
}

fun newtonRaphson(
	initialApproximation: Double,
	function: (Double) -> Double,
	derivative: (Double) -> Double,
	tolerance: Double = -1.0,
	epsilon: Double = Double.MIN_VALUE,
	maximumIterations: Int = 100,
): Double {
	require(maximumIterations > 0) {
		"maximumIterations ($maximumIterations) must be greater than 0"
	}
	require(tolerance != 0.0) {
		"tolerance can't be zero"
	}

	var currentApproximation = initialApproximation

	repeat(maximumIterations) {
		val previousApproximation = currentApproximation
		val derivativeResult = derivative(previousApproximation)

		if (derivativeResult < epsilon || derivativeResult == Double.MIN_VALUE) {
			return Double.NaN
		}
		if (derivativeResult == 0.0) {
			return Double.NaN
		}
		val nextApproximation = previousApproximation - (function(previousApproximation) / derivativeResult)

		val currentRelativeError = abs(1 - previousApproximation / nextApproximation)
		if (tolerance > 0 && currentRelativeError < tolerance) {
			return nextApproximation
		}

		currentApproximation = nextApproximation
	}

	if (tolerance < 0) {
		return currentApproximation
	}

	throw RuntimeException("Method did not converge within $maximumIterations iterations")
}

fun newtonRaphson(
	initialApproximation: Double,
	tolerance: Double = -1.0,
	maximumIterations: Int = 100,
	delta: Double = 0.0000001,
	epsilon: Double = Double.MIN_VALUE,
	function: (Double) -> Double,
): Double {
	require(maximumIterations > 0) {
		"maximumIterations ($maximumIterations) must be greater than 0"
	}
	require(tolerance != 0.0) {
		"tolerance can't be zero"
	}

	var currentApproximation = initialApproximation

	repeat(maximumIterations) {
		val previousApproximation = currentApproximation
		val functionResult = function(previousApproximation)
		val derivativeResult = (function(previousApproximation + delta) - functionResult) / delta

		if (derivativeResult < epsilon || derivativeResult == Double.MIN_VALUE) {
			return Double.NaN
		}
		if (derivativeResult == 0.0) {
			return Double.NaN
		}
		val nextApproximation = previousApproximation - (function(previousApproximation) / derivativeResult)

		val currentRelativeError = abs(1 - previousApproximation / nextApproximation)
		if (tolerance > 0 && currentRelativeError < tolerance) {
			return nextApproximation
		}

		currentApproximation = nextApproximation
	}

	if (tolerance < 0) {
		return currentApproximation
	}

	throw RuntimeException("Method did not converge within $maximumIterations iterations")
}