test_knot_and_data_generation.py

from generate_data import * # this automatically imports the generate_knots and all variables generated by it
import pandas as pd
import pints
import sys
import time as tm
import random
import csv
import os
import pickle as pkl

# main
if __name__ == '__main__':
    # set up variables for the simulation
    tlim = [300, 14899]
    times = np.linspace(*tlim, tlim[-1] - tlim[0], endpoint=False)
    del tlim
    model_name = 'HH'
    snr_db = 20
    state_names = ['a', 'r']  # how many states we have in the model that we are fitting
    ####################################################################################################################
    ## from here there should be no user-changable variables!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    ####################################################################################################################
    # load prortocols
    load_protocols
    # generate the segments with B-spline knots and intialise the betas for splines
    jump_indeces, times_roi, voltage_roi, knots_roi, spline_order = generate_knots(times)
    nSegments = len(jump_indeces[:-1])
    print('Inner optimisation is split into ' + str(nSegments) + ' segments based on protocol steps.')
    nBetas = (len(knots_roi[0]) - spline_order - 1) # number of B-spline coefficients per state per segment
    nBsplineCoeffs = nBetas * len(state_names) # number of B-spline coefficients per segment for all states
    init_betas_roi = nSegments * [0.5 * np.ones(nBsplineCoeffs)]
    print('Number of B-spline coeffs per segment: ' + str(nBsplineCoeffs) + '.')
    ####################################################################################################################
    # generate synthetic data
    if model_name.lower() not in available_models:
        raise ValueError(f'Unknown model name: {model_name}. Available models are: {available_models}')
    elif model_name.lower() == 'hh':
        thetas_true = thetas_hh_baseline
    elif model_name.lower() == 'kemp':
        thetas_true = thetas_kemp
    g = thetas_true[-1] # the last parameter is the conductance - get it as a separate variable just in case
    solution, current_model = generate_synthetic_data(model_name, thetas_true, times)
    states_true = solution.sol(times)
    snr = 10 ** (snr_db / 10)
    current_true = current_model(times, solution, thetas_true, snr=snr)
    states_roi, states_known_roi, current_roi = split_generated_data_into_segments(solution, current_true, jump_indeces, times)
    print('Produced synthetic data for the ' + model_name +' model based on the pre-loaded voltage protocol.')
    ####################################################################################################################
    ## make indexing of B-spline coeffs generalisable for a set number of hidden states
    nBsplineCoeffs = nBetas * len(state_names)
    ## create a list of indeces to insert first B-spline coeffs for each segment
    indeces_to_add = [0]
    for iState in range(1, len(state_names)):
        indeces_to_add.append((nBetas - 1) * iState)
    ## create a list of indeces to drop from the B-spline coeff sets for each segment
    indeces_to_drop = [0]
    for iState in range(1, len(state_names)):
        indeces_to_drop.append(int(nBetas * iState))
    upper_bound_beta = 0.99