ODE_model_code_standalone.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "515356d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from findiff import FinDiff, coefficients\n",
    "from scipy.linalg import solve, det, expm\n",
    "from math import isclose\n",
    "import pandas as pd\n",
    "from scipy.signal import savgol_filter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "d205e57f",
   "metadata": {},
   "outputs": [],
   "source": [
    "############ ODE MODEL #################\n",
    "\n",
    "### FOR FINITE DIFFERENCE ###\n",
    "\n",
    "def floor_even(num):\n",
    "    return np.floor(num/2)*2\n",
    "\n",
    "def floor_odd(num):\n",
    "    if num%2==1:\n",
    "        return num\n",
    "    else:\n",
    "        return num-1\n",
    "\n",
    "def highest_possible_N(time_series):\n",
    "    number = 2/3 * len(time_series) -2/3\n",
    "\n",
    "    # Due to floating point errors causing np.floor to return the wrong number,\n",
    "    # (for instance number=65.9999 when it is actually =66), we use isclose \n",
    "    if isclose(number,np.ceil(number), abs_tol=1e-8):\n",
    "        return int(round(number))\n",
    "    else:\n",
    "        return int(np.floor(number))\n",
    "\n",
    "def highest_possible_accuracies(time_series, N):\n",
    "    acc_list = []\n",
    "    for deriv in range(1,N+1):\n",
    "        # Different conditions for odd and even derivatives:\n",
    "        if deriv%2==1:\n",
    "            number = 2/3*(len(time_series))+1-deriv\n",
    "            # Handle floating point error:\n",
    "            if isclose(number,np.ceil(number), abs_tol=1e-8):\n",
    "                actual_number=round(number)\n",
    "                acc_list.append(int(floor_even(actual_number)))\n",
    "            else:\n",
    "                acc_list.append(int(floor_even(number)))\n",
    "        else:\n",
    "            number = 2/3*(len(time_series))+4/3-deriv\n",
    "            # Handle floating point error:\n",
    "            if isclose(number,np.ceil(number), abs_tol=1e-8):\n",
    "                actual_number=round(number)\n",
    "                acc_list.append(int(floor_even(actual_number)))\n",
    "            else:\n",
    "                acc_list.append(int(floor_even(number)))\n",
    "                \n",
    "    return acc_list\n",
    "\n",
    "\n",
    "def points_to_accuracy(time_series, p_to_use, N):\n",
    "    print(\"p_to_use: \", p_to_use)\n",
    "    # n+m points used in a findiff approx (order + accuracy)\n",
    "    # Accuracy has to be an even number.\n",
    "    \n",
    "    accuracy_list = []\n",
    "    \n",
    "    if p_to_use == 'few':\n",
    "        accuracy_list = [2]*N\n",
    "        \n",
    "    if p_to_use == 'all':\n",
    "        accuracy_list = highest_possible_accuracies(time_series, N)\n",
    "        \n",
    "    if '%' in p_to_use:\n",
    "        percent = float(p_to_use.split(\"%\")[0])*0.01\n",
    "        ts_len = len(time_series)\n",
    "        num_points = int(np.round(ts_len*percent))\n",
    "        \n",
    "        for n in range(1,N+1):\n",
    "            if n>num_points-2:\n",
    "                m = 2\n",
    "                accuracy_list.append(m)\n",
    "            elif n%2==1:\n",
    "                m = int(floor_even(num_points-n))\n",
    "                while ((n + m - 1)/2 - 1 + n + m) > ts_len:\n",
    "                    m -= 2\n",
    "                accuracy_list.append(m)\n",
    "            elif n%2==0:     \n",
    "                m = int(floor_even(num_points-n))\n",
    "                while ((n + m - 1)/2 - 1 - 1 + n + m) > ts_len:\n",
    "                    m -= 2\n",
    "                accuracy_list.append(m)    \n",
    "    return accuracy_list\n",
    "\n",
    "def derivative_approximations(h, time_series, N, accuracy_list):\n",
    "    fin_diff = [time_series] # the 0-th derivative is the original time series\n",
    "    #np.set_printoptions(precision=2)\n",
    "    for k in range(1,N+1):\n",
    "        derivative = FinDiff(0, h, k, acc=accuracy_list[k-1])\n",
    "        fin_diff.append(derivative(time_series))\n",
    "        \n",
    "        mat_repr_deriv_approx = derivative.matrix(time_series.shape)\n",
    "        #print(f\"FINITE DIFF. MATRIX REPRESENTATION. deriv={k}, acc={accuracy_list[k-1]}\")\n",
    "        #print(mat_repr_deriv_approx.toarray())\n",
    "        #print(\"\")\n",
    "        \n",
    "        \n",
    "    return np.asarray(fin_diff)\n",
    "\n",
    "### FOR DETERMINING COEFFICIENTS ANALYTICALLY ###\n",
    "\n",
    "def coeff_analytical_matrix_vector(fin_diff, coeff_to_lock):\n",
    "    N = np.shape(fin_diff)[0]-1 # Highest order derivative\n",
    "        \n",
    "    # We will return A and b, as we later wish to solve the matrix equation Ax=b, where the solution x \n",
    "    # will contain all the coefficients except the locked one, i.e. c_k: [a, c_0, ..., c_k-1, c_k+1, ..., c_N]:\n",
    "    A = np.zeros((N+1,N+1)) # LHS MATRIX\n",
    "    b = np.zeros((N+1, 1))  # RHS vector\n",
    "        \n",
    "        \n",
    "    # Two cases. If coeff_to_lock is 'a':\n",
    "    if coeff_to_lock == 'a':\n",
    "        temp_matrix = np.copy(fin_diff)\n",
    "        for m in range(N+1):\n",
    "            b[m,0] = -1*np.sum(temp_matrix[m])\n",
    "            A[m,m] = np.sum(temp_matrix[m]*temp_matrix[m]) \n",
    "            for l in range(m+1, N+1):\n",
    "                A_ml = np.sum(temp_matrix[l]*temp_matrix[m])\n",
    "                A[m,l] = A_ml\n",
    "                A[l,m] = A_ml \n",
    "        \n",
    "        # In this case, the final solution x to Ax=b, will look like x=[c_0, c_1, ..., c_k, ..., c_N]\n",
    "        \n",
    "    # coeff_to_lock is c_k:\n",
    "    else:\n",
    "        coeff_to_lock_index = int(coeff_to_lock.split(\"_\")[1]) # e.g. c_0 corresponds to row 0 in fin_diff\n",
    "        \n",
    "        # Code is nicer if we replace row corresponding to locked \n",
    "        # coefficient with ones\n",
    "        dyk_i = np.copy(fin_diff[coeff_to_lock_index])\n",
    "        temp_matrix = np.delete(fin_diff, coeff_to_lock_index, axis=0)\n",
    "        temp_matrix = np.insert(temp_matrix, 0, np.ones_like(dyk_i), axis=0)\n",
    "\n",
    "        for m in range(N+1):\n",
    "            b[m,0] = -1*np.sum(dyk_i*temp_matrix[m])\n",
    "            A[m,m] = np.sum(temp_matrix[m]*temp_matrix[m]) \n",
    "            for l in range(m+1, N+1):\n",
    "                A_ml = np.sum(temp_matrix[l]*temp_matrix[m])\n",
    "                A[m,l] = A_ml\n",
    "                A[l,m] = A_ml \n",
    "                \n",
    "        # In this case, the final solution x to Ax=b will look like x=[a, c_0, ..., c_k-1, c_k+1, ..., c_N]\n",
    "    \n",
    "    return A, b\n",
    "\n",
    "'''\n",
    "# Prevent singular matrix\n",
    "def check_if_deriv_is_zero(fin_diff, tolerance=1e-6):\n",
    "    # We return the index of the first derivative which is practically zero (we know all the following ones will also be zero).\n",
    "    sum_of_rows = np.sum(abs(fin_diff), axis=1)\n",
    "    print(\"Sum of each row in fin_diff: \\n\", sum_of_rows)\n",
    "    for row_number, row_sum in enumerate(sum_of_rows):\n",
    "        if abs(row_sum) <= tolerance:\n",
    "            return row_number\n",
    "    return 'All derivatives nonzero'\n",
    "\n",
    "def prev_sing_matrix(fin_diff, coefficient_to_lock):\n",
    "    ### PREVENT SINGULAR MATRIX: This can happen if the time series is sampled from a polynomial with degree < N.\n",
    "    # If any of the derivative approximations is equal to zero, this will give us a singular matrix below.\n",
    "    # If the derivs after a certain order are zero, we set the coefficients for these derivs to zero.\n",
    "    # Then we only work with fin_diff up to and including the final nonzero deriv approx.   \n",
    "    print(\"PREVENT SINGULAR MATRIX\")\n",
    "    print(\"Check if any of the derivatives are 0: \\n\")\n",
    "    first_index_deriv_zero = check_if_deriv_is_zero(fin_diff, tolerance=1e-6) # Need a tolerance due to floating point error.\n",
    "    print(\"\")\n",
    "    if first_index_deriv_zero!='All derivatives nonzero':\n",
    "        if first_index_deriv_zero == 0:\n",
    "            print(\"The time series is equal to zero.\")\n",
    "            return\n",
    "        \n",
    "        if coefficient_to_lock != 'a' and int(coefficient_to_lock.split(\"_\")[1])>=first_index_deriv_zero:\n",
    "            print(\"The selected coefficient to lock corresponds to a derivative that is equal to zero.\")\n",
    "            print(f\"Select a coefficient c_k to lock with k <= {first_index_deriv_zero-1}, as all derivatives of order\")\n",
    "            print(\"higher than this are equal to zero.\")\n",
    "            return\n",
    "        \n",
    "        fin_diff = fin_diff[:first_index_deriv_zero]\n",
    "        \n",
    "        print(f\"The highest order derivative to include was given as: {N}.\")\n",
    "        print(f\"However, the derivatives of order >= {first_index_deriv_zero} are equal to zero.\")\n",
    "        print(\"The returned list of coefficients will therefore only include the coefficients up to and including:\") \n",
    "        print(f\"c_{first_index_deriv_zero-1} \\n\")\n",
    "    else:\n",
    "        print(\"findiff does not contain rows that are equal to zero.\")\n",
    "    return fin_diff\n",
    "\n",
    "# Prevent linearly dependent matrix\n",
    "def determine_lin_dep_columns(A, tolerance=1e-6):\n",
    "    low = 0\n",
    "    high = A.shape[1]-1\n",
    "    mid = 0\n",
    "    \n",
    "    # We do a binary search until we get the index of the final column to include, whcih corresponds with the \"high\" variable\n",
    "    if abs(det(A))<tolerance:\n",
    "        while low <= high:\n",
    "            mid = (high+low)//2\n",
    "            \n",
    "            if abs(det(A[:mid, :mid]))<tolerance:\n",
    "                high = mid-1\n",
    "            else:\n",
    "                low = mid+1\n",
    "                \n",
    "        return high\n",
    "    else:\n",
    "        return (\"Matrix is linearly independent\")\n",
    "    \n",
    "def prev_lin_dep_matrix(A, b):\n",
    "    ### PREVENT LINEARLY DEPENDENT MATRIX: This can happen if a row in fin_diff is a multiple of \n",
    "    # an earlier row in fin_diff. I.e. what happens when we have a time series sampled from exp(t),\n",
    "    # then every dervative is the same. We solve this by finding the largest submatrix in A s.t.\n",
    "    # A is linearly independent. We do this by checking if the determinant is smaller or bigger\n",
    "    # than some tolerance.\n",
    "    print(\"CHECK IF MATRIX IS LINEARLY INDEPENDENT\")\n",
    "    final_col_to_include = determine_lin_dep_columns(np.copy(A), tolerance=1e-4)\n",
    "    if final_col_to_include != \"Matrix is linearly independent\":\n",
    "        print(\"The matrix was not linearly independent.\")\n",
    "        print(\"We extract the largest submatrix starting in the upper left corner that's linearly independent.\")\n",
    "        print(f\"The final column from A we include is the column with index: {final_col_to_include-1}.\")\n",
    "        print(f\"This gives the matrix A[:{final_col_to_include},:{final_col_to_include}]:\")\n",
    "        \n",
    "        A = A[:final_col_to_include, :final_col_to_include]\n",
    "        b = b[:final_col_to_include]\n",
    "        \n",
    "        print(\"A: \\n\", A)\n",
    "        print(\"b: \\n\", b)\n",
    "        \n",
    "        print(f\"Since the matrix {final_col_to_include} rows/columns corresponds to having originally\")\n",
    "        print(f\"chosen N={final_col_to_include-1} instead, the returned list of coefficients\")\n",
    "        print(f\"will therefore only include coefficients up to and including: c_{final_col_to_include-1}\")\n",
    "    else:\n",
    "        print(final_col_to_include)\n",
    "    return A, b\n",
    "'''\n",
    "\n",
    "\n",
    "def insert_locked_coeff(solved_coeff, coefficient_to_lock):\n",
    "    if coefficient_to_lock == 'a':\n",
    "        solved_all_coeff = np.insert(solved_coeff, 0, 1, axis=0) \n",
    "        # Here, solved_all_coeff looks like [1, c_0, ..., c_k, ..., c_N]\n",
    "    else: \n",
    "        coeff_to_lock_index = int(coefficient_to_lock.split(\"_\")[1])+1 # +1 since 'a' is at index 0.\n",
    "        solved_all_coeff = np.insert(solved_coeff, coeff_to_lock_index, 1, axis=0)\n",
    "        # Here, if c_k is the locked coeff, solved_all_coeff looks like [a, c_0, ..., c_k-1, 1, c_k+1, ..., c_N].\n",
    "    return solved_all_coeff\n",
    "        \n",
    "# Find the coefficients solving the minimization problem\n",
    "def solved_all_coeffs_analytical(time_points, time_series, N, acc_p_to_use, coefficient_to_lock='c_0', accuracy_list=None):\n",
    "    time_points = np.array(time_points)\n",
    "    time_series = np.array(time_series)\n",
    "    \n",
    "    ### Check if chosen N is possible:\n",
    "    highest_pos_N = highest_possible_N(time_series)\n",
    "    if N<=highest_pos_N:\n",
    "        print(f\"The chosen highest order derivative is N={N}, which is less than or equal\")\n",
    "        print(f\"to the highest possible derivative, N_highest = {highest_pos_N}\")\n",
    "    else:\n",
    "        print(f\"The chosen highest order derivative is N={N}, which is greater than\")\n",
    "        print(f\"the highest possible derivative, N_highest = {highest_pos_N}\")\n",
    "        return\n",
    "    \n",
    "    # Check if coefficient to lock is valid\n",
    "    if coefficient_to_lock != 'a' and coefficient_to_lock[:2]!='c_':\n",
    "        print(\"coefficient_to_lock must be either 'a' or 'c_k', for 0 <= k <= N\")\n",
    "        return\n",
    "    \n",
    "    if coefficient_to_lock[:2]=='c_' and (int(coefficient_to_lock.split(\"_\")[1])>N or int(coefficient_to_lock.split(\"_\")[1])<0):\n",
    "        print(\"c_k: k must be between 0 and N\")\n",
    "        return\n",
    "    \n",
    "    ### Create a list with the highest possible accuracy for each derivative approx:\n",
    "    if accuracy_list == None:\n",
    "        accuracy_list = points_to_accuracy(time_series, acc_p_to_use, N)\n",
    "    \n",
    "    print(\"\\n List containing order of accuracy for the derivative approximations.\")\n",
    "    print(\"Index i corresponds to the accuracy of the derivative of order i+1: \\n\", accuracy_list, \"\\n\")\n",
    "    \n",
    "    ### Create an array with all derivative approximations. \n",
    "    ### The row corresponds to the derivative order, and column corresponds to the time point.\n",
    "    ### Ex: value at (2, 1), is the approximation of the second derivative of y at time t_1, that is y_1''.\n",
    "    fin_diff = derivative_approximations(time_points[1]-time_points[0], time_series, N, accuracy_list)  \n",
    "    print(\"\\n fin_diff: \\n\", fin_diff, \"\\n\")\n",
    "    ###  fin_diff - numpy array with  shape (N+1, n+1): \n",
    "    ###    [ [y_0     ,   y_1      ,   ...   ,    y_n     ],\n",
    "    ###      [y_0'    ,   y_1'     ,   ...   ,    y_n'    ],\n",
    "    ###      [y_0''   ,   y_1''    ,   ...   ,    y_n''   ],\n",
    "    ###                     ...,\n",
    "    ###      [y_0^(N) ,   y_1^(N)  ,   ...   ,    y_n^(N) ]  ] \n",
    "    \n",
    "    \n",
    "    #fin_diff = prev_sing_matrix(fin_diff, coefficient_to_lock)\n",
    "    \n",
    "    print(\"CONSTRUCT MATRIX FOR MATRIX EQUATION\")\n",
    "    A, b = coeff_analytical_matrix_vector(fin_diff, coefficient_to_lock) \n",
    "    print(\"A: \\n\", A)\n",
    "    print(\"b: \\n\", b)\n",
    "    print(\"\")    \n",
    "    #A, b = prev_lin_dep_matrix(A,b)\n",
    "    print(\"det of A: \", det(A))\n",
    "    \n",
    "    solved_coeff = solve(A, b, assume_a='sym') # Column vector. We use the convention [\"a\",\"c_0\", \"c_1\", \"c_2\", ..., \"c_N\"]^T\n",
    "    \n",
    "    # The solution to the matrix equation Ax = b does not contain the locked coefficient, so we have to add it back.\n",
    "    # Since this is the coefficient that's \"divided out\", it has value 1:\n",
    "    solved_all_coeff = insert_locked_coeff(solved_coeff, coefficient_to_lock)\n",
    "    \n",
    "    return solved_all_coeff, fin_diff\n",
    "\n",
    "### FOR CALCULATING THE EXPLICIT SOLUTION\n",
    "\n",
    "# Find the matrix and vector for converting ODE to system of first order equations\n",
    "def coeff_list_to_system_matrix(all_coeffs, tolerance):\n",
    "    # all_coeffs = np.array([[a], [1], [c_1], [c_2], ..., [c_N]]) where c_0 is the locked coeff.\n",
    "    \n",
    "    # As we will divide by c_N, we must ensure that the last coeff in the list is nonzero.\n",
    "    # This is OK, becaue if the array ends with c_{k+1}, ..., c_{N} that are (approx) equal to zero\n",
    "    # we in practice have an ODE of order k, so we can just drop these last terms. \n",
    "    while abs(all_coeffs[-1,0])<tolerance:\n",
    "        all_coeffs = all_coeffs[:-1]\n",
    "    \n",
    "    print(\"remove highest order coeff(s) if below threshold: \\n\", all_coeffs)\n",
    "    \n",
    "    # The order of the ODE determines the shape of the matrix.\n",
    "    ode_order = all_coeffs.shape[0]-2 # -2 because we have the constant term and the y^(0) coeff.\n",
    "    print(\"If any coefficients were removed, the effective ODE order will change. The ODE order is now: \", ode_order)\n",
    "    if ode_order == 0:\n",
    "        C = None\n",
    "        A = None\n",
    "        return C, A\n",
    "        \n",
    "    # Create matrix C\n",
    "    C = -1*np.eye(ode_order, k=1)\n",
    "    bottom_row = all_coeffs[1:-1,0].T/all_coeffs[-1,0]\n",
    "    print(\"bottom_row: \", bottom_row)\n",
    "    C[-1, :] = bottom_row \n",
    "    print(\"C: \\n\", C)\n",
    "    # Create vector A\n",
    "    A = np.zeros((ode_order,1))\n",
    "    A[-1,0]=all_coeffs[0,0]\n",
    "    print(\"A: \\n\", A)\n",
    "    return C, A\n",
    "\n",
    "def initial_cond(fin_diff, percentage_of_interval, where_to_start, N):\n",
    "    num_points = floor_odd(int(np.floor(fin_diff.shape[1]*0.01*float(percentage_of_interval.split('%')[0]))))\n",
    "    print(\"num_points: \", num_points)\n",
    "    t0_index = 0\n",
    "    \n",
    "    if where_to_start == 'start':\n",
    "        t0_index = num_points//2 + 1\n",
    "    \n",
    "    if where_to_start == 'middle':\n",
    "        t0_index = fin_diff.shape[1]//2\n",
    "        \n",
    "    if where_to_start == 'end':\n",
    "        t0_index = fin_diff.shape[1] - (num_points//2 + 1)\n",
    "    \n",
    "    print(\"t0_index: \", t0_index)\n",
    "    \n",
    "    y_initial = [np.mean(fin_diff[0,t0_index-num_points//2:t0_index+num_points//2+1])]\n",
    "    for n in range(1,N):\n",
    "        if n%2==1:\n",
    "            new_num_points = num_points-(2+n)\n",
    "        else:\n",
    "            new_num_points = num_points-(3+n)\n",
    "        \n",
    "        y_initial.append(np.mean(fin_diff[n,t0_index-new_num_points//2:t0_index+new_num_points//2+1]))\n",
    "    \n",
    "    return np.reshape(y_initial,(-1,1)), t0_index\n",
    "\n",
    "# Calculate explicit solution to system of first order equations\n",
    "def Cinverse_negA(C,A):\n",
    "    print(\"CINVERSENEGA\")\n",
    "    print(\"C: \\n\", C)\n",
    "    print(\"A: \\n\", A)\n",
    "    print(\"np.linalg.inv(C): \\n\", np.linalg.inv(C))\n",
    "    return np.matmul(np.linalg.inv(C), -A)\n",
    "\n",
    "def first_order_system_sol(t, t_0, y_0, C, Cinverse_negA):\n",
    "    mat_exp = expm(C*(t_0-t))\n",
    "    #if t<0.02:\n",
    "    #    print(\"C: \", C)\n",
    "    #    print(\"mat_exp: \", mat_exp)\n",
    "    #print(\"Cinverse_negA: \", Cinverse_negA)\n",
    "    return np.matmul(mat_exp, y_0-Cinverse_negA)+Cinverse_negA\n",
    "\n",
    "\n",
    "### Code for the model which puts the above together ###\n",
    "\n",
    "def model(time, time_series, N, acc_points_to_use, init_points_to_use, init_pos, coeff_lock='c_0', acc_list=None):\n",
    "    print(\"\")\n",
    "    print(\"\")\n",
    "    print(\"!!!!!!! START !!!!!!!!!\")\n",
    "    num_points_to_use = [1]*N\n",
    "    initial_time = 't_0'\n",
    "    tolerance = 1e-6\n",
    "    \n",
    "    all_coeffs, fin_diff = solved_all_coeffs_analytical(time, time_series, N, acc_points_to_use, coefficient_to_lock=coeff_lock, accuracy_list=acc_list)\n",
    "    print(\"fin_diff: \\n\", fin_diff)\n",
    "    print(\"!!!!!!!!!!!!!\")\n",
    "    print(\"all_coeffs: \\n\", all_coeffs)\n",
    "\n",
    "    C, A = coeff_list_to_system_matrix(all_coeffs, tolerance)\n",
    "    \n",
    "    try:\n",
    "        print(\"Start of try\")\n",
    "        num_derivs = C.shape[0]+1\n",
    "        print(\"num_derivs: \", num_derivs)\n",
    "        fin_diff = fin_diff[:num_derivs]\n",
    "        print(\"If the effective ODE order is different, we correspondingly remove rows from fin_diff: \\n\", fin_diff)\n",
    "        #y_init = initial_conditions(fin_diff, num_points_to_use, initial_time)\n",
    "        y_init, t0_index = initial_cond(fin_diff, init_points_to_use, init_pos, num_derivs-1)\n",
    "        print(\"initial conditions: \\n\", y_init)\n",
    "        #t_0 = time[int(initial_time.split(\"_\")[1])]\n",
    "        t_0 = time[t0_index]\n",
    "        print(\"init time: \", t_0)\n",
    "\n",
    "        Cinv_negA = Cinverse_negA(C,A)\n",
    "        print(\"Cinv_negA: \\n\", Cinv_negA)\n",
    "        y_sol = []\n",
    "\n",
    "        for t in time:\n",
    "            y_sol.append(first_order_system_sol(t, t_0, y_init, C, Cinv_negA))\n",
    "\n",
    "        y_sol_forward = y_sol.copy()\n",
    "        for t in (time+time[-1]):\n",
    "            y_sol_forward.append(first_order_system_sol(t, t_0, y_init, C, Cinv_negA))\n",
    "        print(\"End of try\")\n",
    "        \n",
    "        return np.array(y_sol), np.array(y_sol_forward), (t0_index, t_0, y_init, C, Cinv_negA), all_coeffs\n",
    "\n",
    "    # If the coefficients in front of all the derivatives are 0:\n",
    "    except:\n",
    "        print(\"In except\")\n",
    "        y_sol = np.ones_like(time)*(-1)*all_coeffs[0,0]    \n",
    "        y_sol = np.reshape(y_sol, (-1,1,1))\n",
    "        y_sol_forward = y_sol.copy()\n",
    "        \n",
    "        return np.array(y_sol), np.array(y_sol_forward), (None, None, None, None, None), all_coeffs\n",
    "\n",
    "\n",
    "### Function that loops over and locks each coefficient, comparing which gives the best result ###\n",
    "\n",
    "def best_c_j_to_lock(time, true_ts, N, acc_points_to_use, init_points_to_use, init_pos):\n",
    "    y_sols_c_j_locked = []\n",
    "    rel_errors = []\n",
    "    model_params_list = []\n",
    "    all_coeffs_list = []\n",
    "    \n",
    "    for j in range(N+1):\n",
    "        y_sol, y_sol_forward, model_params, all_coeffs = model(time, true_ts, N, acc_points_to_use, init_points_to_use, init_pos, coeff_lock=f'c_{j}')\n",
    "        y_sols_c_j_locked.append(y_sol)\n",
    "        rel_errors.append(relative_error(y_sol[:,0,0], true_ts))\n",
    "        model_params_list.append(model_params)\n",
    "        all_coeffs_list.append(all_coeffs)\n",
    "    best_j = rel_errors.index(np.nanmin(np.array(rel_errors)))\n",
    "    print(\"\")\n",
    "    print(\"\")\n",
    "    print(\"Relative errors: \\n\", rel_errors)\n",
    "    print(f\"Best c_j to lock: c_{best_j}\", )\n",
    "    return y_sols_c_j_locked[best_j], rel_errors[best_j], model_params_list[best_j], best_j, all_coeffs_list[best_j]\n",
    "\n",
    "### Function that loops through N (up to and including highest_N), and compares the results ###\n",
    "def best_N(time, true_ts, ts_name, highest_N, plot_best_of_each_N, init_points_to_use, init_pos, save_figure=False, acc_points_to_use='few'):\n",
    "    \n",
    "    best_sols_dict = {}\n",
    "    best_errors_dict = {}\n",
    "    best_model_params_dict = {}\n",
    "    best_locked_j_dict = {}\n",
    "    best_all_coeffs_dict = {}\n",
    "    for N in range(1,highest_N+1):\n",
    "        y_best, best_rel_error, best_model_params, best_locked_j, best_all_coeffs = best_c_j_to_lock(time, true_ts, N, acc_points_to_use,  init_points_to_use, init_pos)\n",
    "        best_sols_dict[N] = y_best\n",
    "        best_errors_dict[N] = best_rel_error\n",
    "        best_model_params_dict[N] = best_model_params\n",
    "        best_locked_j_dict[N] = best_locked_j\n",
    "        best_all_coeffs_dict[N] = best_all_coeffs\n",
    "        \n",
    "    best_sols_and_errors = best_errors_dict.items()\n",
    "    \n",
    "    best_N = min(best_errors_dict, key=best_errors_dict.get)\n",
    "    \n",
    "    best_y = best_sols_dict[best_N]\n",
    "    best_err = best_errors_dict[best_N]\n",
    "    best_model_params = best_model_params_dict[best_N]\n",
    "    \n",
    "    #print(\"best_sols_dict: \\n\", best_sols_dict)\n",
    "    print(\"\")\n",
    "    print(\"\")\n",
    "    print(\"!!\")\n",
    "    print(\"best_locked_j_dict: \", best_locked_j_dict)\n",
    "    print(\"best_errors_dict: \\n\", best_errors_dict)\n",
    "    print(\"best_all_coeffs_dict: \\n\", best_all_coeffs_dict)\n",
    "    print(\"best_N: \", best_N)\n",
    "    print(\"best_locked_c_j: \", best_locked_j_dict[best_N])\n",
    "    if plot_best_of_each_N:\n",
    "        for N, best_y in best_sols_dict.items():\n",
    "            plt.title(ts_name)         \n",
    "            #plt.plot(time, el_price_2019_data[:len(time)], color=\"gray\",alpha=0.5, label=\"Original\")\n",
    "            plt.plot(time, true_ts, label='True')\n",
    "            plt.plot(time, best_y[:,0,0], label=f'Pred. N={N}. Locked c_{best_locked_j_dict[N]}.\\nInit. cond: {init_points_to_use}, {init_pos}')\n",
    "            plt.legend()\n",
    "            \n",
    "            if save_figure:\n",
    "                if N==best_N:\n",
    "                    plt.savefig(savefigs_loc+ts_name+f\"__BEST__N_{N}_best_locked_{best_locked_j_dict[N]}_init_{init_points_to_use.split('%')[0]}_{init_pos}__TRAINING__rel_error_{best_errors_dict[N]}.png\")\n",
    "                else:\n",
    "                    plt.savefig(savefigs_loc+\"not_best_figures\\\\\"+ts_name+f\"__NOT_BEST__N_{N}_best_locked_{best_locked_j_dict[N]}_init_{init_points_to_use.split('%')[0]}_{init_pos}__TRAINING__rel_error_{best_errors_dict[N]}.png\")       \n",
    "            plt.show()\n",
    "            print(\"Training, relative error: \", best_errors_dict[N])\n",
    "            print(\"ODE Coefficients: \\n\", best_all_coeffs_dict[N])\n",
    "            \n",
    "    return best_y, best_err, best_model_params, best_N, best_locked_j_dict[best_N]\n",
    "\n",
    "### Function that calculates and plots function given user-specified coefficients and initial values ###\n",
    "def pre_chosen_coeffs_model(chosen_coeffs, chosen_init_conds, time_interval, t0=0):\n",
    "    C, A = coeff_list_to_system_matrix(chosen_coeffs, 1e-10)\n",
    "    Cinv_negA = Cinverse_negA(C,A)\n",
    "    y_pred_chosen=[]\n",
    "    for t in time_interval:\n",
    "        y_pred_chosen.append(first_order_system_sol(t, t0, chosen_init_conds, C, Cinv_negA))\n",
    "        \n",
    "    plt.plot(time_interval, np.array(y_pred_chosen)[:,0,0])\n",
    "    #plt.savefig(savefigs_loc+f\"pre_chosen_coeffs_sol_{time_interval[-1]}.png\")\n",
    "    plt.show()\n",
    "    return np.array(y_pred_chosen)[:,0,0]\n",
    "\n",
    "### Function that takes in model parameters, and gives the model prediction forward in time ###\n",
    "def model_params_testing(time, best_model_params, true, title, forward_in_time, best_N, best_locked_j, save_figure = False):\n",
    "    sol_list = []\n",
    "\n",
    "    t0_index = best_model_params[0]\n",
    "    t0 = best_model_params[1]\n",
    "    y_init = best_model_params[2]\n",
    "    C = best_model_params[3]\n",
    "    Cinv_negA = best_model_params[4]\n",
    "\n",
    "    for t in time:\n",
    "        sol = first_order_system_sol(t, t0, y_init, C, Cinv_negA)\n",
    "        sol_list.append(sol)\n",
    "        \n",
    "    \n",
    "    rel_error=0 \n",
    "    if forward_in_time:\n",
    "        middle_index = time.shape[0]//2\n",
    "        rel_error = relative_error(np.array(sol_list)[middle_index:,0,0],true[middle_index:])\n",
    "        print(\"Testing (forward in time), Relative error (on future time points): \", rel_error)\n",
    "        \n",
    "        #plt.plot(time, tesla_2021_2022_data[:len(time)], color=\"gray\",alpha=0.5, label=\"Original\")\n",
    "        #plt.plot(time, true, label='Smoothed')\n",
    "        plt.plot(time, np.array(sol_list)[:,0,0], label='Pred. data')\n",
    "        plt.title(title)\n",
    "        plt.legend()\n",
    "        if save_figure:\n",
    "            plt.savefig(savefigs_loc+title+f\"__best_N_{best_N}_best_locked_{best_locked_j}__FIT__rel_error_{rel_error}.png\")\n",
    "    else:\n",
    "        rel_error = relative_error(np.array(sol_list)[:,0,0],true[:])\n",
    "        print(\"Testing (same interval), Relative error: \", rel_error)\n",
    "        \n",
    "        #plt.plot(time, tesla_2021_data, color=\"gray\", alpha=0.5, label=\"Original\")\n",
    "        #plt.plot(time, true, label='Smoothed')\n",
    "        plt.plot(time, np.array(sol_list)[:,0,0], label='Pred. data')\n",
    "        plt.title(title)\n",
    "        plt.legend()\n",
    "        if save_figure:\n",
    "            plt.savefig(savefigs_loc+title+f\"__best_N_{best_N}_best_locked_{best_locked_j}__TESTING__rel_error_{rel_error}.png\")\n",
    "    \n",
    "    plt.show()\n",
    "    return"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}