itrails 0.1.0__py3-none-any.whl

itrails/__init__.py

@@ -0,0 +1,4 @@
+try:
+    from ._version import version as __version__  # type: ignore
+except ImportError:
+    __version__ = "unknown"

itrails/_version.py

@@ -0,0 +1,21 @@
+# file generated by setuptools-scm
+# don't change, don't track in version control
+
+__all__ = ["__version__", "__version_tuple__", "version", "version_tuple"]
+
+TYPE_CHECKING = False
+if TYPE_CHECKING:
+    from typing import Tuple
+    from typing import Union
+
+    VERSION_TUPLE = Tuple[Union[int, str], ...]
+else:
+    VERSION_TUPLE = object
+
+version: str
+__version__: str
+__version_tuple__: VERSION_TUPLE
+version_tuple: VERSION_TUPLE
+
+__version__ = version = '0.1.0'
+__version_tuple__ = version_tuple = (0, 1, 0)

itrails/combine_states.py

@@ -0,0 +1,140 @@
+import numba as nb
+import numpy as np
+
+
+def combine_states(
+    state_dict_1, state_dict_2, state_dict_sum, final_probs_1, final_probs_2
+):
+    """
+    Function that combines two dictionaries of individual states and their final probabilities into a single dictionary of combined states and starting probabilities.
+
+    :param state_dict_1: Dictionary of states and indices for the first CTMC, can be 1 sequence or 2 sequence CTMC.
+    :type state_dict_1: Numba Dictionary of Key: Tuple of int64 and Value: int64.
+    :param state_dict_2: Dictionary of states and indices for the second CTMC, can only be 1 sequence CTMC.
+    :type state_dict_2: Numba Dictionary of Key: Tuple(int64, int64) and Value: int64
+    :param state_dict_sum: Dictionary of states and indices for the combined CTMC.
+    :type state_dict_sum: Numba Dictionary of Key: Tuple of int64 and Value: int64.
+    :param final_probs_1: Array of final probabilities for a single key of the first CTMC.
+    :type final_probs_1: Array of type: float64[:, :]
+    :param final_probs_2: Array of final probabilities for the only key of the second CTMC.
+    :type final_probs_2: Array of type: float64[:, :]
+    :return: Array of starting probabilities for each combined state in the combined CTMC.
+    :rtype: Array of type: float64[:, :]
+    """
+    len_array = len(list(state_dict_1.keys())[0]) + len(list(state_dict_2.keys())[0])
+    init_comb_dict = {}
+    init_comb_all = np.zeros((len(state_dict_sum.keys())), dtype=np.float64)
+    for key_1, index_1 in state_dict_1.items():
+        left_1 = key_1[: len(key_1) // 2]
+        right_1 = key_1[len(key_1) // 2 :]
+        for key_2, index_2 in state_dict_2.items():
+            left_2 = key_2[: len(key_2) // 2]
+            right_2 = key_2[len(key_2) // 2 :]
+            comb_state = np.zeros((len_array), dtype=np.int64)
+            used_values_1 = {}
+            used_values_2 = {}
+            current_value = 1
+            index = 0
+            for value in left_1:
+                if value in used_values_1:
+                    comb_state[index] = used_values_1[value]
+                else:
+                    used_values_1[value] = current_value
+                    comb_state[index] = current_value
+                    current_value += 1
+                index += 1
+
+            for value in left_2:
+                if value in used_values_2:
+                    comb_state[index] = used_values_2[value]
+                else:
+                    used_values_2[value] = current_value
+                    comb_state[index] = current_value
+                    current_value += 1
+                index += 1
+
+            for value in right_1:
+                if value in used_values_1:
+                    comb_state[index] = used_values_1[value]
+                else:
+                    used_values_1[value] = current_value
+                    comb_state[index] = current_value
+                    current_value += 1
+                index += 1
+
+            for value in right_2:
+                if value in used_values_2:
+                    comb_state[index] = used_values_2[value]
+                else:
+                    used_values_2[value] = current_value
+                    comb_state[index] = current_value
+                    current_value += 1
+                index += 1
+
+            init_comb_dict[tuple(comb_state)] = (
+                final_probs_1[0, index_1] * final_probs_2[0, index_2]
+            )
+    for state in init_comb_dict.keys():
+        index_AB = state_dict_sum[state]
+        init_comb_all[index_AB] = init_comb_dict[state]
+    return init_comb_all
+
+
+def combine_states_wrapper(
+    state_dict_1,
+    state_dict_2,
+    state_dict_sum,
+    final_probs_1,
+    final_probs_2,
+):
+    """
+    Wrapper function that combines dictionaries of states and their final probabilities into a single dictionary of combined states and starting probabilities.
+
+    :param state_dict_1: Dictionary of states and indices for the first CTMC, can be 1 sequence or 2 sequence CTMC.
+    :type state_dict_1: Numba Dictionary of Key: Tuple of int64 and Value: int64.
+    :param state_dict_2: Dictionary of states and indices for the second CTMC, can only be 1 sequence CTMC.
+    :type state_dict_2: Numba Dictionary of Key: Tuple(int64, int64) and Value: int64.
+    :param state_dict_sum: Dictionary of states and indices for the combined CTMC.
+    :type state_dict_sum: Numba Dictionary of Key: Tuple of int64 and Value: int64.
+    :param final_probs_1: Final probability dictionary for the first CTMC.
+    :type final_probs_1: Numba Dictionary of Key: UniTuple(nb.types.UniTuple(int64, 3), 2) and Value: float64[:, :].
+    :param final_probs_2: Final probability dictionary for the second CTMC.
+    :type final_probs_2: Numba Dictionary of Key: UniTuple(nb.types.UniTuple(int64, 3), 2) and Value: float64[:, :].
+    :raises NotImplementedError: Not implemented for more than 2 species in state_dict_1 or and more than 1 species in state_dict_2.
+    :raises Exception: Fallback if invalid format in final_probs_1 or final_probs_2.
+    :return: Dictionary of combined states and starting probabilities for each state.
+    :rtype: Numba Dictionary of Key: UniTuple(nb.types.UniTuple(int64, 3), 2) and Value: float64[:, :].
+    """
+    pi_dict = nb.typed.Dict.empty(
+        key_type=nb.types.UniTuple(nb.types.UniTuple(nb.types.int64, 3), 2),
+        value_type=nb.types.float64[:, :],
+    )
+
+    start_placeholder = ((-1, -1, -1), (-1, -1, -1))
+
+    if len(final_probs_1.keys()) > 1 and len(final_probs_2.keys()) > 1:
+        raise NotImplementedError
+
+    elif len(final_probs_1.keys()) > 1 and len(final_probs_2.keys()) == 1:
+        prob2 = final_probs_2[start_placeholder]
+        for path, prob1 in final_probs_1.items():
+            pi_vector_combined = combine_states(
+                state_dict_1, state_dict_2, state_dict_sum, prob1, prob2
+            )
+
+            pi_dict[path] = pi_vector_combined.reshape(1, -1)
+
+        return pi_dict
+
+    elif len(final_probs_1.keys()) == 1 and len(final_probs_2.keys()) == 1:
+
+        prob1 = final_probs_1[start_placeholder]
+        prob2 = final_probs_2[start_placeholder]
+        pi_vector_combined = combine_states(
+            state_dict_1, state_dict_2, state_dict_sum, prob1, prob2
+        )
+        pi_dict[start_placeholder] = pi_vector_combined.reshape(1, -1)
+        return pi_dict
+
+    else:
+        raise Exception("Invalid final_probs_1 or final_probs_2")

itrails/cutpoints.py

@@ -0,0 +1,65 @@
+import numpy as np
+from scipy.stats import expon, truncexpon
+
+
+def cutpoints_AB(n_int_AB, t_AB, coal_AB):
+    """
+    This function returns the cutpoints for the
+    intervals for the two-sequence CTMC. The cutpoints
+    will be defined by the quantiles of a truncated
+    exponential distribution.
+
+    :param n_int_AB: Number of intervals in the two-sequence CTMC.
+    :type n_int_AB: int
+    :param t_AB: Total time interval of the two-sequence CTMC
+    :type t_AB: float
+    :param coal_AB: coalescent rate of the two-sequence CTMC.
+    :type coal_AB: float
+    :return: cut_AB
+    :rtype: np.array
+    """
+    quantiles_AB = np.array(list(range(n_int_AB + 1))) / n_int_AB
+    lower, upper, scale = 0, t_AB, 1 / coal_AB
+    cut_AB = truncexpon.ppf(
+        quantiles_AB, b=(upper - lower) / scale, loc=lower, scale=scale
+    )
+    return cut_AB
+
+
+def cutpoints_ABC(n_int_ABC, coal_ABC):
+    """
+    This function returns the cutpoints for the
+    intervals for the three-sequence CTMC. The cutpoints
+    will be defined by the quantiles of an exponential
+    distribution.
+
+    :param n_int_ABC: Number of intervals in the three-sequence CTMC.
+    :type n_int_ABC: int
+    :param coal_ABC: coalescent rate of the three-sequence CTMC.
+    :type coal_ABC: float
+    :return: cut_ABC
+    :rtype: np.array
+    """
+    quantiles_AB = np.array(list(range(n_int_ABC + 1))) / n_int_ABC
+    cut_ABC = expon.ppf(quantiles_AB, scale=1 / coal_ABC)
+    return cut_ABC
+
+
+def get_times(cut, intervals):
+    """
+    Returns a list of time differences for each specified interval.
+
+    This function computes the duration for each interval defined by the given indices.
+    It does so by subtracting the earlier cutpoint from the subsequent one based on the indices
+    provided in the intervals list.
+
+    :param cut: List of ordered cutpoints.
+    :type cut: list[float]
+    :param intervals: Ordered indices of cutpoints that define the intervals.
+    :type intervals: list[int]
+    :return: List of time differences between consecutive cutpoints as defined by intervals.
+    :rtype: list[float]
+    """
+    return [
+        cut[intervals[i + 1]] - cut[intervals[i]] for i in range(len(intervals) - 1)
+    ]

itrails/deepest_ti.py

@@ -0,0 +1,256 @@
+import numpy as np
+
+
+def deep_identify(
+    current,
+    absorbing_state,
+    omega_nonrev_counts,
+    inverted_omega_nonrev_counts,
+    path,
+    all_paths_dict,
+    by_l=-1,
+    by_r=-1,
+):
+    """
+    Recursive function that identifies, in each iteration, the possible paths that can be taken from the current state to the final state. The function is called recursively until the final state is reached. The function stores the paths in the paths_array and all_paths_array arrays.
+
+    :param current: Current omega state in the recursion
+    :type current: Tuple of int
+    :param absorbing_state: Absorbing state
+    :type absorbing_state: Tuple of int
+    :param omega_nonrev_counts: Dictionary containing the number of non-reversible coalescents (value) for each omega state (key)
+    :type omega_nonrev_counts: Numba typed Dict
+    :param inverted_omega_nonrev_counts: Dictionary containing the omega states (value) for each number of non-reversible coalescents (key)
+    :type inverted_omega_nonrev_counts: Numba typed Dict
+    :param path: Current path in the recursion
+    :type path: Numpy array
+    :param all_paths_dict: Dictionary that recursively gets filled up with all the paths
+    :type all_paths_dict: Numpy array
+    :param by_l: Current omega left subpath, defaults to -1 (initial placeholder value)
+    :type by_l: int, optional
+    :param by_r: Current omega right subpath, defaults to -1 (initial placeholder value)
+    :type by_r: int, optional
+    """
+    # Calculate differences
+    diff_l = omega_nonrev_counts[absorbing_state[0]] - omega_nonrev_counts[current[0]]
+    diff_r = omega_nonrev_counts[absorbing_state[1]] - omega_nonrev_counts[current[1]]
+
+    # Termination condition
+    if diff_l <= 1 and diff_r <= 1:
+        key = (by_l, by_r)
+        if key not in all_paths_dict:
+            all_paths_dict[key] = []
+        # Create a copy of the current path
+        path_copy = [tuple(p) for p in path]
+        all_paths_dict[key].append(path_copy)
+        return
+
+    start_l = omega_nonrev_counts[current[0]]
+    start_r = omega_nonrev_counts[current[1]]
+    end_l = omega_nonrev_counts[absorbing_state[0]]
+    end_r = omega_nonrev_counts[absorbing_state[1]]
+
+    # Explore next states for left
+    if start_l < end_l:
+        next_states_l = inverted_omega_nonrev_counts[start_l + 1]
+        for left in next_states_l:
+            new_state = (left, current[1])
+            new_by_l = (
+                by_l
+                if by_l != -1
+                else (
+                    left
+                    if omega_nonrev_counts[left] == 1 and start_l + 1 != end_l
+                    else -1
+                )
+            )
+            path.append(new_state)
+            deep_identify(
+                new_state,
+                absorbing_state,
+                omega_nonrev_counts,
+                inverted_omega_nonrev_counts,
+                path,
+                all_paths_dict,
+                new_by_l,
+                by_r,
+            )
+            path.pop()
+
+    # Explore next states for right
+    if start_r < end_r:
+        next_states_r = inverted_omega_nonrev_counts[start_r + 1]
+        for right in next_states_r:
+            new_state = (current[0], right)
+            new_by_r = (
+                by_r
+                if by_r != -1
+                else (
+                    right
+                    if omega_nonrev_counts[right] == 1 and start_r + 1 != end_r
+                    else -1
+                )
+            )
+            path.append(new_state)
+            deep_identify(
+                new_state,
+                absorbing_state,
+                omega_nonrev_counts,
+                inverted_omega_nonrev_counts,
+                path,
+                all_paths_dict,
+                by_l,
+                new_by_r,
+            )
+            path.pop()
+
+    # Explore next states for both left and right
+    if start_l < end_l and start_r < end_r:
+        next_states_l = inverted_omega_nonrev_counts[start_l + 1]
+        next_states_r = inverted_omega_nonrev_counts[start_r + 1]
+        for left in next_states_l:
+            for right in next_states_r:
+                if omega_nonrev_counts[right] > start_r:
+                    new_state = (left, right)
+                    new_by_l = (
+                        by_l
+                        if by_l != -1
+                        else (
+                            left
+                            if omega_nonrev_counts[left] == 1 and start_l + 1 != end_l
+                            else -1
+                        )
+                    )
+                    new_by_r = (
+                        by_r
+                        if by_r != -1
+                        else (
+                            right
+                            if omega_nonrev_counts[right] == 1 and start_r + 1 != end_r
+                            else -1
+                        )
+                    )
+                    path.append(new_state)
+                    deep_identify(
+                        new_state,
+                        absorbing_state,
+                        omega_nonrev_counts,
+                        inverted_omega_nonrev_counts,
+                        path,
+                        all_paths_dict,
+                        new_by_l,
+                        new_by_r,
+                    )
+                    path.pop()
+
+
+def deep_identify_wrapper(
+    omega_init,
+    absorbing_state,
+    omega_nonrev_counts,
+    inverted_omega_nonrev_counts,
+    path_to_convert,
+):
+    """
+    Wrapper function for the deep_identify function. This function initializes the arrays and dictionaries needed for the recursion and calls the deep_identify function. In the end it returns the transformed keys and the paths.
+
+    :param omega_init: Initial omega state
+    :type omega_init: Tuple of int
+    :param absorbing_state: Absorbing state
+    :type absorbing_state: Tuple of int
+    :param omega_nonrev_counts: Dictionary containing the number of non-reversible coalescents (value) for each omega state (key)
+    :type omega_nonrev_counts: Numba typed Dict
+    :param inverted_omega_nonrev_counts: Dictionary containing the omega states (value) for each number of non-reversible coalescents (key)
+    :type inverted_omega_nonrev_counts: Numba typed Dict
+    :return: Resulting keys and paths
+    :rtype: Tuple(Array, Array, Array, int)
+    """
+
+    all_paths_dict = {}
+    path = [omega_init]
+    deep_identify(
+        omega_init,
+        absorbing_state,
+        omega_nonrev_counts,
+        inverted_omega_nonrev_counts,
+        path,
+        all_paths_dict,
+    )
+
+    # Convert dictionary to keys_array and paths_array
+    keys_array = np.array(list(all_paths_dict.keys()))
+    keys_array_final = np.zeros((len(keys_array), 6))
+    path_to_convert_array = np.array(path_to_convert)
+    flattened_array = np.hstack(path_to_convert_array.ravel())
+    coded_dict = {3: 1, 5: 2, 6: 3}
+    for i, key in enumerate(keys_array):
+        keys_array_final[i] = flattened_array
+        if key[0] != -1 and keys_array_final[i][0] == -1:
+            keys_array_final[i][0] = coded_dict[key[0]]
+        if key[1] != -1 and keys_array_final[i][3] == -1:
+            keys_array_final[i][3] = coded_dict[key[1]]
+
+    # Find the maximum number of paths and subpaths for padding
+    max_paths = max(len(paths) for paths in all_paths_dict.values())
+    max_subpaths = max(
+        max(len(subpath) for subpath in paths) for paths in all_paths_dict.values()
+    )
+
+    # Initialize paths_array with zeros
+    paths_array = np.zeros((len(all_paths_dict), max_paths, max_subpaths, 2))
+
+    # Initialize path_lengths_array to store the length of each path
+    path_lengths_array = np.zeros((len(all_paths_dict), max_paths))
+
+    # Fill paths_array and path_lengths_array
+    for i, (key, paths) in enumerate(all_paths_dict.items()):
+        for j, subpath in enumerate(paths):
+            path_lengths_array[i, j] = len(subpath)  # Store the length of the path
+            for k, point in enumerate(subpath):
+                paths_array[i, j, k] = point
+
+    return keys_array_final, paths_array, path_lengths_array, max_subpaths
+
+
+def deepest_ti(trans_mat_noabs, omega_dict_noabs, path):
+    """
+    This function calculated the integral of matrix exponentials with an infinite time limit.
+
+    :param trans_mat_noabs: Transition matrix without absorbing states
+    :type trans_mat_noabs: Numpy array
+    :param omega_dict_noabs: Omega dictionary without absorbing states
+    :type omega_dict_noabs: Numpy array
+    :param path: Path of omega states
+    :type path: Numpy array
+    :return: Result of the integral of the series of multiplying matrix exponentials.
+    :rtype: Numpy array
+    """
+    steps = len(path) - 1
+    n = trans_mat_noabs.shape[0]
+    if steps == 1:
+        C_mat = trans_mat_noabs
+
+    elif steps > 1:
+        C_mat = np.zeros((n * steps, n * steps))
+        C_mat[0:n, 0:n] = trans_mat_noabs
+        for idx in range(1, steps):
+
+            sub_om_init = (path[idx - 1, 0], path[idx - 1, 1])
+            sub_om_fin = (path[idx, 0], path[idx, 1])
+            A_mat = (
+                np.diag(omega_dict_noabs[sub_om_init].astype(np.float64))
+                @ trans_mat_noabs
+                @ np.diag(omega_dict_noabs[sub_om_fin].astype(np.float64))
+            )
+            C_mat[n * idx : n * (idx + 1), n * idx : n * (idx + 1)] = trans_mat_noabs
+            C_mat[n * (idx - 1) : n * idx, n * idx : n * (idx + 1)] = A_mat
+
+    sub_om_init = (path[-2, 0], path[-2, 1])
+    sub_om_fin = (path[-1, 0], path[-1, 1])
+    A_mat = np.ascontiguousarray(
+        np.diag(omega_dict_noabs[sub_om_init].astype(np.float64))
+        @ trans_mat_noabs
+        @ np.diag(omega_dict_noabs[sub_om_fin].astype(np.float64))
+    )
+    result = (-np.linalg.inv(C_mat))[:n, -n:] @ A_mat
+    return result

itrails/examples/example_config.yaml

@@ -0,0 +1,21 @@
+# Example configuration for itrails-optimize
+fixed_parameters:
+  mu: 2e-8
+
+optimized_parameters: # [starting, min, max]
+  N_AB: [50000, 5000, 500000]
+  N_ABC: [50000, 5000, 500000]
+  t_1: [240000, 24000, 2400000]
+  t_2: [40000, 4000, 400000]
+  t_3: [800000, 80000, 8000000]
+  t_upper: [745069.3855665945, 74506.93855665945, 7450693.855665945]
+  r: [1e-8, 1e-9, 1e-7]
+
+settings:
+  input_maf: # Path to the MAF alignment file (overwritten by console argument)
+  output_name: # Path to the output directory (overwritten by console argument)
+  n_cpu: 64
+  method: "Nelder-Mead"
+  species_list: ["hg38", "panTro5", "gorGor5", "ponAbe2"]
+  n_int_AB: 3
+  n_int_ABC: 3

itrails/expm.py

@@ -0,0 +1,166 @@
+from __future__ import division, print_function
+
+import math
+
+import numba as nb
+import numpy as np
+
+
+@nb.jit(nopython=True, parallel=False, fastmath=True)
+def expm(A):
+    """
+    Calculates matrix exponential of a square matrix A.
+    Adapted from https://github.com/michael-hartmann/expm/blob/master/python/expm.py
+    Algorithm 10.20 from unctions of Matrices: Theory and Computation, Nicholas J. Higham, 2008
+    :param A: square matrix
+    :type A: np.array
+    :return: matrix exponential of A
+    :rtype: np.array
+    """
+    theta3 = 1.5e-2
+    theta5 = 2.5e-1
+    theta7 = 9.5e-1
+    theta9 = 2.1e0
+    theta13 = 5.4e0
+    # calculate the norm of A
+    norm = np.linalg.norm(A, ord=1)
+
+    if norm < theta3:
+        dtype = A.dtype
+        dim, dim = A.shape
+        b = [120, 60, 12, 1]
+
+        U = b[1] * np.eye(dim, dtype=dtype)
+        V = b[0] * np.eye(dim, dtype=dtype)
+
+        A2 = A @ A
+        A2n = np.eye(dim, dtype=dtype)
+
+        # evaluate (10.33)
+        for i in range(1, 3 // 2 + 1):
+            A2n = A2n @ A2
+            U += b[2 * i + 1] * A2n
+            V += b[2 * i] * A2n
+
+        U = A @ U
+        return np.linalg.solve(V - U, V + U)
+
+    elif norm < theta5:
+        dtype = A.dtype
+        dim, dim = A.shape
+        b = [30240, 15120, 3360, 420, 30, 1]
+
+        U = b[1] * np.eye(dim, dtype=dtype)
+        V = b[0] * np.eye(dim, dtype=dtype)
+
+        A2 = A @ A
+        A2n = np.eye(dim, dtype=dtype)
+
+        # evaluate (10.33)
+        for i in range(1, 5 // 2 + 1):
+            A2n = A2n @ A2
+            U += b[2 * i + 1] * A2n
+            V += b[2 * i] * A2n
+
+        U = A @ U
+        return np.linalg.solve(V - U, V + U)
+
+    elif norm < theta7:
+        dtype = A.dtype
+        dim, dim = A.shape
+        b = [17297280, 8648640, 1995840, 277200, 25200, 1512, 56, 1]
+
+        U = b[1] * np.eye(dim, dtype=dtype)
+        V = b[0] * np.eye(dim, dtype=dtype)
+
+        A2 = A @ A
+        A2n = np.eye(dim, dtype=dtype)
+
+        # evaluate (10.33)
+        for i in range(1, 7 // 2 + 1):
+            A2n = A2n @ A2
+            U += b[2 * i + 1] * A2n
+            V += b[2 * i] * A2n
+
+        U = A @ U
+        return np.linalg.solve(V - U, V + U)
+
+    elif norm < theta9:
+        dtype = A.dtype
+        dim, dim = A.shape
+        b = [
+            17643225600,
+            8821612800,
+            2075673600,
+            302702400,
+            30270240,
+            2162160,
+            110880,
+            3960,
+            90,
+            1,
+        ]
+
+        U = b[1] * np.eye(dim, dtype=dtype)
+        V = b[0] * np.eye(dim, dtype=dtype)
+
+        A2 = A @ A
+        A2n = np.eye(dim, dtype=dtype)
+
+        # evaluate (10.33)
+        for i in range(1, 9 // 2 + 1):
+            A2n = A2n @ A2
+            U += b[2 * i + 1] * A2n
+            V += b[2 * i] * A2n
+
+        U = A @ U
+        return np.linalg.solve(V - U, V + U)
+
+    else:
+
+        # algorithm 10.20, from line 7
+        dim, dim = A.shape
+        b = [
+            64764752532480000,
+            32382376266240000,
+            7771770303897600,
+            1187353796428800,
+            129060195264000,
+            10559470521600,
+            670442572800,
+            33522128640,
+            1323241920,
+            40840800,
+            960960,
+            16380,
+            182,
+            1,
+        ]
+
+        s = max(0, int(math.ceil(math.log(norm / theta13) / math.log(2))))
+        if s > 0:
+            A /= 2**s
+
+        Id = np.eye(dim)
+        A2 = A @ A
+        A4 = A2 @ A2
+        A6 = A2 @ A4
+
+        U = A @ (
+            (A6 @ (b[13] * A6 + b[11] * A4 + b[9] * A2))
+            + b[7] * A6
+            + b[5] * A4
+            + b[3] * A2
+            + b[1] * Id
+        )
+
+        V = (
+            A6 @ (b[12] * A6 + b[10] * A4 + b[8] * A2)
+            + b[6] * A6
+            + b[4] * A4
+            + b[2] * A2
+            + b[0] * Id
+        )
+        # added this
+        r13 = np.ascontiguousarray(np.linalg.solve(V - U, V + U))
+        return np.linalg.matrix_power(r13, 2**s)