Needs "books.txt" in the resources directory. Try experimenting with weights on line 136

layout_generator.kt

package name.anton3.layout

import java.util.*
import kotlin.collections.ArrayList
import kotlin.math.abs

class Dummy

fun readResource(name: String): String {
    return Dummy::class.java.classLoader.getResource(name)!!.readText()
}

fun makeLayoutConversions(lines: List<String>): Pair<Map<Char, Byte>, Map<Byte, Char>> {
    val layout = lines.joinToString("").toCharArray()
    val layoutToIndex: Map<Char, Byte> = layout.withIndex().associateBy({ it.value }, { it.index.toByte() })
    val layoutFromIndex: Map<Byte, Char> = layout.withIndex().associateBy({ it.index.toByte() }, { it.value })
    return layoutToIndex to layoutFromIndex
}

@Suppress("NOTHING_TO_INLINE")
inline fun bigramIndex(letter1: Byte, letter2: Byte, nLetters: Int): Int {
    return letter1.toInt() * nLetters + letter2
}

@Suppress("NOTHING_TO_INLINE")
inline fun indexToBigram(index: Int, nLetters: Int): Pair<Byte, Byte> {
    return (index / nLetters).toByte() to (index % nLetters).toByte()
}

fun makeSameFingerBigramIndices(fingers: ByteArray): IntArray {
    val result = ArrayList<Int>()

    for (finger in 0..9) {
        val fingerIndices = fingers.withIndex().filter { it.value == finger.toByte() }.map { it.index.toByte() }

        for (idx1 in fingerIndices) {
            for (idx2 in fingerIndices) {
                if (idx1 != idx2) result.add(bigramIndex(idx1, idx2, fingers.size))
            }
        }
    }

    return result.toIntArray()
}

fun makeSameHandBigramIndices(fingers: ByteArray): IntArray {
    val result = ArrayList<Int>()

    for (hand in 0..1) {
        val handIndices = fingers.withIndex().filter { it.value / 5 == hand }.map { it.index.toByte() }

        for (idx1 in handIndices) {
            for (idx2 in handIndices) {
                if (idx1 != idx2) result.add(bigramIndex(idx1, idx2, fingers.size))
            }
        }
    }

    return result.toIntArray()
}

fun makeRollsBigramIndices(fingers: ByteArray): Pair<IntArray, IntArray> {
    val inward = ArrayList<Int>()
    val outward = ArrayList<Int>()

    for ((key1, finger1) in fingers.withIndex()) {
        for ((key2, finger2) in fingers.withIndex()) {

            val leftyComparison = (when {
                finger1 / 5 != finger2 / 5 -> 0
                finger1 / 5 == 0 -> 1
                else -> -1
            }) * finger1.compareTo(finger2)

            val rollType = when {
                leftyComparison < 0 -> inward
                leftyComparison > 0 -> outward
                else -> null
            }.takeIf { abs(key1 / 10 - key2 / 10) <= 1 }

            rollType?.add(bigramIndex(key1.toByte(), key2.toByte(), fingers.size))
        }
    }

    return inward.toIntArray() to outward.toIntArray()
}

private const val NULL_INDEX: Byte = -1

fun letterAndBigramWeights(
    corpus: String,
    lowerToIndex: Map<Char, Byte>,
    upperToIndex: Map<Char, Byte>
): Pair<DoubleArray, DoubleArray> {
    val nLetters = lowerToIndex.size

    val letterCounts = IntArray(nLetters)
    val bigramCounts = IntArray(nLetters * nLetters)
    var previousLetter: Byte = NULL_INDEX

    for (c in corpus) {
        if (!c.isLetter()) continue

        val idx = lowerToIndex[c] ?: upperToIndex[c]
        if (idx == null) {
            previousLetter = NULL_INDEX
            continue
        }

        letterCounts[idx.toInt()] += 1
        if (previousLetter != NULL_INDEX) {
            bigramCounts[bigramIndex(previousLetter, idx, nLetters)] += 1
        }

        previousLetter = idx
    }

    val letterTotal = letterCounts.sumByDouble { it.toDouble() }
    val bigramTotal = bigramCounts.sumByDouble { it.toDouble() }

    val letterWeights = DoubleArray(letterCounts.size) { letterCounts[it].toDouble() / letterTotal }
    val bigramWeights = DoubleArray(bigramCounts.size) { bigramCounts[it].toDouble() / bigramTotal }

    return letterWeights to bigramWeights
}

data class Energy(
    val key: Double,
    val sameFinger: Double,
    val sameHand: Double,
    val inwardRolls: Double,
    val outwardRolls: Double
) : Comparable<Energy> {

    @Suppress("NOTHING_TO_INLINE")
    inline fun total(): Double = key * 2.0 + sameHand * 6.0 + sameFinger * 3.0 - inwardRolls * 4.0 - outwardRolls * 2.0

    override fun compareTo(other: Energy): Int = this.total().compareTo(other.total())

    override fun toString(): String =
        String.format(Locale.ENGLISH, "%.4f (key: %.4f, same-finger: %.4f, same-hand: %.4f, inward: %.4f, outward: %.4f)", total(), key, sameFinger, sameHand, inwardRolls, outwardRolls)
}

class OptimizationWeights(
    val keyWeights: DoubleArray,
    val bigramWeights: DoubleArray,
    val fingerPenalties: DoubleArray,
    val sameFingerBigramIndices: IntArray,
    val sameHandBigramIndices: IntArray,
    val inwardRollsBigramIndices: IntArray,
    val outwardRollsBigramIndices: IntArray
)

@Suppress("NOTHING_TO_INLINE")
private inline fun sumBigramFreq(mapping: ByteArray, bigramIndices: IntArray, bigramWeights: DoubleArray): Double {
    val nLetters = mapping.size
    return bigramIndices.sumByDouble { fromBigramIndex ->
        val (from1, from2) = indexToBigram(fromBigramIndex, nLetters)
        val toBigramIndex = bigramIndex(mapping[from1.toInt()], mapping[from2.toInt()], nLetters)
        bigramWeights[toBigramIndex]
    }
}

fun evaluateEnergy(mapping: ByteArray, weights: OptimizationWeights): Energy {
    val keyEnergy = mapping.indices.sumByDouble { from ->
        weights.fingerPenalties[from] * weights.keyWeights[mapping[from].toInt()]
    }
    val sameFingerFreq = sumBigramFreq(mapping, weights.sameFingerBigramIndices, weights.bigramWeights)
    val sameHandFreq = sumBigramFreq(mapping, weights.sameHandBigramIndices, weights.bigramWeights)
    val inwardRollFreq = sumBigramFreq(mapping, weights.inwardRollsBigramIndices, weights.bigramWeights)
    val outwardRollFreq = sumBigramFreq(mapping, weights.outwardRollsBigramIndices, weights.bigramWeights)
    return Energy(
        key = keyEnergy,
        sameFinger = sameFingerFreq,
        sameHand = sameHandFreq,
        inwardRolls = inwardRollFreq,
        outwardRolls = outwardRollFreq
    )
}

fun layoutAsString(
    mapping: ByteArray,
    lowerFromIndex: Map<Byte, Char>,
    upperFromIndex: Map<Byte, Char>
): String {
    return mapping.asList().chunked(10) { letters ->
        letters.joinToString(" ") { letter ->
            "${lowerFromIndex[letter]}"
        }
    }.joinToString("\n")
}

data class StartingLayout(
    val lower: List<String>,
    val upper: List<String>,
    val movable: List<String>
) {
    init {
        require(lower.map { it.length } == upper.map { it.length })
        require(lower.map { it.length } == movable.map { it.length })
        require(movable.joinToString("").all { it in "01" })
    }
}

@Suppress("SpellCheckingInspection")
val colemakNoDiagLayout = StartingLayout(
    lower = listOf(
        "qwfp,.luyg",
        "arstdhneio",
        "zxcv/;mbkj"
    ),
    upper = listOf(
        "QWFP<>LUYG",
        "ARSTDHNEIO",
        "ZXCV?:MBKJ"
    ),
    movable = listOf(
        "1111001111",
        "1111111111",
        "1111001111"
    )
)

@Suppress("SpellCheckingInspection")
val colemakLayout: StartingLayout = StartingLayout(
    lower = listOf(
        "qwfpgjluy;",
        "arstdhneio",
        "zxcvbkm,./"
    ),
    upper = listOf(
        "QWFPGJLUY:",
        "ARSTDHNEIO",
        "ZXCVBKM<>?"
    ),
    movable = listOf(
        "1111111111",
        "1111111111",
        "1111111111"
    )
)

val ergoFingers: List<String> = listOf(
    "0123366789",
    "0123366789",
    "0123366789"
)

val fingerPenaltiesBase: DoubleArray = doubleArrayOf(
    2.7, 2.7, 2.3, 2.3, 2.8, 2.8, 2.3, 2.3, 2.7, 2.7,
    1.0, 0.5, 0.0, 0.0, 2.0, 2.0, 0.0, 0.0, 0.5, 1.0,
    3.5, 3.2, 2.5, 2.5, 3.6, 3.6, 2.5, 2.5, 3.2, 3.5
)

val fingerPenaltiesNoDiag: DoubleArray = doubleArrayOf(
    2.7, 2.7, 2.3, 2.3, 4.5, 4.5, 2.3, 2.3, 2.7, 2.7,
    1.0, 0.5, 0.0, 0.0, 2.5, 2.5, 0.0, 0.0, 0.5, 1.0,
    3.5, 3.2, 2.5, 2.5, 4.5, 4.5, 2.5, 2.5, 3.2, 3.5
)

val fingerPenaltiesNoDiagNoPinkies: DoubleArray = doubleArrayOf(
    2.7, 2.7, 2.3, 2.3, 4.5, 4.5, 2.3, 2.3, 2.7, 2.7,
    2.0, 1.0, 0.0, 0.0, 2.5, 2.5, 0.0, 0.0, 1.0, 2.0,
    3.5, 3.2, 2.5, 2.5, 4.5, 4.5, 2.5, 2.5, 3.2, 3.5
)

@Suppress("NOTHING_TO_INLINE")
inline fun ByteArray.swap(idx1: Int, idx2: Int) {
    val temp = this[idx1]
    this[idx1] = this[idx2]
    this[idx2] = temp
}

fun keySwaps(mapping: ByteArray, movable: BooleanArray): Sequence<Unit> = sequence {
    for (i in (0 until 30).shuffled()) {
        for (j in (0 until i).shuffled()) {
            if (!movable[i] || !movable[j]) continue
            mapping.swap(i, j)
            yield(Unit)
        }
    }

    for (i1 in (0 until 30).shuffled()) {
        for (j1 in (0 until i1).shuffled()) {
            if (!movable[i1] || !movable[j1]) continue

            for (i2 in (0 until 30).shuffled()) {
                for (j2 in (0 until i2).shuffled()) {
                    if (!movable[i2] || !movable[j2]) continue
                    mapping.swap(i1, j1)
                    mapping.swap(i2, j2)
                    yield(Unit)
                }
            }
        }
    }

    for (i in (0 until 10).shuffled()) {
        for (j in (0 until i).shuffled()) {
            if (!movable[i + 0] || !movable[j + 0] ||
                !movable[i + 10] || !movable[j + 10] ||
                !movable[i + 20] || !movable[j + 20]
            ) continue

            mapping.swap(i + 0, j + 0)
            mapping.swap(i + 10, j + 10)
            mapping.swap(i + 20, j + 20)
            yield(Unit)
        }
    }

    for (i1 in (0 until 10).shuffled()) {
        for (j1 in (0 until i1).shuffled()) {
            if (!movable[i1 + 0] || !movable[j1 + 0] ||
                !movable[i1 + 10] || !movable[j1 + 10] ||
                !movable[i1 + 20] || !movable[j1 + 20]
            ) continue

            for (i2 in (0 until 10).shuffled()) {
                for (j2 in (0 until i2).shuffled()) {
                    if (!movable[i2 + 0] || !movable[j2 + 0] ||
                        !movable[i2 + 10] || !movable[j2 + 10] ||
                        !movable[i2 + 20] || !movable[j2 + 20]
                    ) continue

                    mapping.swap(i1 + 0, j1 + 0)
                    mapping.swap(i1 + 10, j1 + 10)
                    mapping.swap(i1 + 20, j1 + 20)
                    mapping.swap(i2 + 0, j2 + 0)
                    mapping.swap(i2 + 10, j2 + 10)
                    mapping.swap(i2 + 20, j2 + 20)
                    yield(Unit)
                }
            }
        }
    }

    for (i1 in (0 until 30).shuffled()) {
        for (j1 in (0 until i1).shuffled()) {
            if (!movable[i1] || !movable[j1]) continue

            for (i2 in (0..j1).shuffled()) {
                for (j2 in (0 until i2).shuffled()) {
                    if (!movable[i2] || !movable[j2]) continue

                    for (i3 in (0..j2).shuffled()) {
                        for (j3 in (0 until 30).shuffled()) {
                            if (!movable[i3] || !movable[j3]) continue
                            mapping.swap(i1, j1)
                            mapping.swap(i2, j2)
                            mapping.swap(i3, j3)
                            yield(Unit)
                        }
                    }
                }
            }
        }
    }
}

fun main() {
    val layout = colemakNoDiagLayout

    val (layoutLowerToIndex, layoutLowerFromIndex) = makeLayoutConversions(layout.lower)
    val (layoutUpperToIndex, layoutUpperFromIndex) = makeLayoutConversions(layout.upper)

    check(layoutLowerToIndex.size == layoutUpperToIndex.size)
    val nLetters = layoutLowerToIndex.size

    val movable: BooleanArray = layout.movable.joinToString("").map { it == '1' }.toBooleanArray()

    val fingers: ByteArray = ergoFingers.joinToString("")
        .map { c -> Character.digit(c, 10).also { check(it != -1) }.toByte() }.toByteArray()

    val sameFingerBigramIndices = makeSameFingerBigramIndices(fingers)
    val sameHandBigramIndices = makeSameHandBigramIndices(fingers)
    val (inwardBigramIndices, outwardBigramIndices) = makeRollsBigramIndices(fingers)

    val fingerPenalties = fingerPenaltiesNoDiagNoPinkies

    val corpus = readResource("books.txt")

    val (letterWeights, bigramWeights) = letterAndBigramWeights(corpus, layoutLowerToIndex, layoutUpperToIndex)

    val weights = OptimizationWeights(
        letterWeights,
        bigramWeights,
        fingerPenalties,
        sameFingerBigramIndices,
        sameHandBigramIndices,
        inwardBigramIndices,
        outwardBigramIndices
    )

    val mapping = ByteArray(nLetters) { it.toByte() }
    var energy = evaluateEnergy(mapping, weights)
    var maxEnergy = energy
    val testMapping = ByteArray(nLetters)

    println(layoutAsString(mapping, layoutLowerFromIndex, layoutUpperFromIndex))
    println("Energy: $energy")
    println()

    var iter = 0

    noReset@while (true) {
        mapping.copyInto(testMapping)

        for (ks in keySwaps(testMapping, movable)) {
            ++iter

            val testEnergy = evaluateEnergy(testMapping, weights)

            if (testEnergy < energy) {
                testMapping.copyInto(mapping)
                energy = testEnergy

                if (energy < maxEnergy) {
                    maxEnergy = energy
                    println(layoutAsString(mapping, layoutLowerFromIndex, layoutUpperFromIndex))
                    println("Energy $iter: $energy")
                    println()
                }
                continue@noReset
            }

            mapping.copyInto(testMapping)
        }

        mapping.indices.forEach { mapping[it] = it.toByte() }
        energy = evaluateEnergy(mapping, weights)
        iter = 0
        println("Reset")
    }
}