The Math of Dying in Pathfinder 2e v2

import numpy as np
from numpy.linalg import matrix_power as mp, inv

def roll_die(state, max_state, die, mod):
    new_state = state

    if die == 20: new_state -= 2
    elif (die <= state + mod) or (die == 1): new_state += 2
    elif die >= 10 + state + mod: new_state -= 1
    else: new_state += 1

    if new_state >= max_state: new_state = max_state
    elif new_state <= 0: new_state = 0

    return new_state

def generate_absorbing_matrix(max_state, mod_recovery):
    order_states = [max_state - 1] + [i for i in range(0, max_state-1)] + [max_state]

    array_dying_0 = [0 for _ in range(1, max_state)] + [1, 0]
    array_last_dying = [0 for _ in range(1, max_state + 1)] + [1]

    transition_absorbing_matrix = []

    for state in range(1, max_state):
        state_count = [0] * (max_state + 1)
        for die in range(1,21):
            state_count[order_states[roll_die(state, max_state, die, mod_recovery)]] += 1/20
        transition_absorbing_matrix.append([round(element, 2) for element in state_count])

    transition_absorbing_matrix.append(array_dying_0)
    transition_absorbing_matrix.append(array_last_dying)

    transition_absorbing_matrix = np.array(transition_absorbing_matrix)

    return transition_absorbing_matrix

def generate_matrix_without_absorbing(absorbing_matrix, state_to_remove):
    new_matrix = np.delete(np.delete(absorbing_matrix, state_to_remove, axis = 1), state_to_remove, axis = 0)

    for i in range(pos_dying_0):
        new_matrix[i] = new_matrix[i] * 1/(1-absorbing_matrix[i][state_to_remove])

    return new_matrix

def calculate_relevant_statistics(matrix, n_absorbing):
    n_transient = len(matrix) - n_absorbing
    transient_matrix = matrix[0:n_transient, 0:n_transient]
    transient_absorbing_matrix = matrix[0:n_transient, n_transient:len(matrix)]
    identity = np.identity(len(transient_matrix))

    fundamental_matrix = inv(identity - transient_matrix)

    hadamard_fundamental = fundamental_matrix * fundamental_matrix
    diagonal_fundamental = np.diag(np.diag(fundamental_matrix))

    variance_visits = np.matmul(fundamental_matrix, (2 * diagonal_fundamental - identity)) - hadamard_fundamental

    expected_steps = np.matmul(fundamental_matrix, np.ones(len(transient_matrix)))
    variance_steps = np.matmul((2 * fundamental_matrix - identity),(expected_steps * expected_steps)) - expected_steps

    transient_probabilities = np.matmul(fundamental_matrix - identity, inv(diagonal_fundamental))
    absorbing_probabilities = np.matmul(fundamental_matrix, transient_absorbing_matrix)

    print(expected_steps)
    print()
    print(absorbing_probabilities)
    print()

max_states = [2, 3, 4, 5]
possible_mod_recovery = [0, -1, -4, -5]

for max_state in max_states:
    for mod_recovery in possible_mod_recovery:
        print(f'Max State: {max_state} and Mod Recovery: {mod_recovery}')
        print()
        transition_absorbing_matrix = generate_absorbing_matrix(max_state, mod_recovery)
        # print(transition_absorbing_matrix)
        # print()

        pos_last_dying = max_state
        pos_dying_0 = max_state - 1

        matrix_only_0 = generate_matrix_without_absorbing(transition_absorbing_matrix, pos_last_dying)
        matrix_only_last = generate_matrix_without_absorbing(transition_absorbing_matrix, pos_dying_0)

        calculate_relevant_statistics(transition_absorbing_matrix, 2)
        # calculate_relevant_statistics(matrix_only_0, 1)
        # calculate_relevant_statistics(matrix_only_last, 1)