import numpy as np
import torch
import os
import matplotlib.pyplot as plt
#from torchdiffeq import odeint
from torchdiffeq import odeint_adjoint as odeint
from linear_memory.linear_memory import LinearMemory
import linear_memory.utils as ut
from import_utils import add_path
from scipy.interpolate import interp1d
add_path('pyssa')
import pyssa.ssa as ssa
import pyssa.models.standard_models as sm
add_path('pymbvi')
from pymbvi.models.mjp.autograd_partition_specific_models import SimpleGeneExpression
from pymbvi.util import num_derivative, autograd_jacobian
torch.manual_seed(2007301620)
# get simulation model
pre, post, rates = sm.get_standard_model("simple_gene_expression")
# prepare initial conditions
initial = np.array([0.0, 1.0, 0.0, 0.0])
tspan = np.array([0.0, 3e3])
# set up gene expression model
moment_initial = np.zeros(9)
moment_initial[0:3] = initial[1:4]
import sys
import os
from pathlib import Path
import numpy as np
import torch
import matplotlib.pyplot as plt
from torchdiffeq import odeint_adjoint as odeint
workdir = Path(__file__).resolve().parent.parent.parent
sys.path.append(str(workdir))
from import_utils import add_path
add_path('pyssa')
#pyssa_path = '/Users/christian/Documents/Code/pyssa'
#sys.path.append(pyssa_path)
import pyssa.util as ut
torch.set_default_dtype(torch.float64)
torch.manual_seed(2008181715)
# load data
load_path = os.path.dirname(os.path.realpath(__file__)) + '/data.npz'
data = np.load(load_path)
moment_initial = torch.from_numpy(data['moment_initial'])
rates = data['rates']
A_true = data['A_true']
b_true = data['b_true']
num_samples = data['num_samples']
tspan = data['tspan']
t_plot = data['t_plot']
delta_t = data['delta_t']