# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# +
from pathlib import Path
import shutil

import amici
import amici.petab_import
import amici.petab_objective
from more_itertools import one
import numpy as np
import petab                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                             


import petab_timecourse
from petab_timecourse import (
    import_directory_of_componentwise_files,
)

model_name = 'Cuadros_HealthPlace2020'

simulation_duration = 205  # simulated time
output_density = 10  # number of output timepoints per simulated time unit

petab_path = Path(model_name)

componentwise_timecourses_path = petab_path / 'regimens'

# +
petab_yaml_location = str(one(petab_path.glob('*yaml')))
petab_problem = petab.Problem.from_yaml(str(one(petab_path.glob('*yaml'))))

import copy
original_petab_problem = copy.deepcopy(petab_problem)
# -

timecourse_df, condition_df = import_directory_of_componentwise_files(
    componentwise_timecourses_path,
)
petab_problem = petab.Problem.from_yaml(petab_yaml_location)
petab_problem.condition_df = condition_df
petab_problem.timecourse_df = timecourse_df
petab_problem.sbml_model = petab_timecourse.sbml.add_timecourse_as_events(
    petab_problem,
    timecourse_id=one(petab_problem.timecourse_df.index),
)
petab_problem.condition_df = original_petab_problem.condition_df

amici_model = amici.petab_import.import_petab_problem(petab_problem, force_compile=True)

sim = amici.petab_objective.simulate_petab(petab_problem, amici_model)

from amici.plotting import plotObservableTrajectories
import matplotlib.pyplot as plt
observable_indices=[amici_model.getObservableIds().index('observable_epsilon')]
plotObservableTrajectories(sim['rdatas'][0], model=amici_model, observable_indices=observable_indices)
plt.show()

from libsbml import formulaToL3String
print('trigger\t\tcomponent\tvalue')
for e in petab_problem.sbml_model.getListOfEvents():
    print(formulaToL3String(e.getTrigger().getMath()), end='\t')
    for a in e.getListOfEventAssignments():
        print(a.getVariable(), end='\t\t')
        print(formulaToL3String(a.getMath()))

import pypesto.petab
# Import the problem and compile the model.
_importer = pypesto.petab.PetabImporter.from_yaml(petab_yaml_location)
pypesto_problem = _importer.create_problem(force_compile=True)

# +
from amici.plotting import plotObservableTrajectories
import matplotlib.pyplot as plt

# Simulate and plot the timecourse parameter.
observable_indices=[pypesto_problem.objective.amici_model.getObservableIds().index('observable_epsilon')]

t = np.linspace(0, simulation_duration, simulation_duration*output_density + 1)
pypesto_problem.objective.amici_model.setTimepoints(t)

par_dict = pypesto_problem.objective.par_arr_to_dct(petab_problem.x_nominal)
ignored_pars = set(par_dict.keys()).difference(pypesto_problem.objective.amici_model.getParameterIds())
pypesto_problem.objective.amici_model.setParameterById({
    k: v
    for k, v in par_dict.items()
    if k not in ignored_pars
})
rdata = amici.runAmiciSimulation(
    pypesto_problem.objective.amici_model,
    pypesto_problem.objective.amici_solver,
)
plotObservableTrajectories(
    rdata,
    model=pypesto_problem.objective.amici_model,
    observable_indices=observable_indices,
)
plt.show()

# +
import pypesto.engine
import pypesto.optimize

optimizer = pypesto.optimize.ScipyOptimizer()
engine = pypesto.engine.MultiProcessEngine()

pypesto_result = pypesto.optimize.minimize(
    problem=pypesto_problem,
    optimizer=optimizer,
    n_starts=16,
    engine=engine,
)
# -

import pypesto.visualize
import matplotlib.pyplot as plt
# Plot optimization information.
pypesto.visualize.waterfall(pypesto_result)
pypesto.visualize.parameters(pypesto_result)
plt.show()

# Simulate and plot the MLE
'''
from amici.plotting import plotStateTrajectories, plotObservableTrajectories
import matplotlib.pyplot as plt
t = np.linspace(0, simulation_duration, simulation_duration*output_density + 1)
pypesto_problem.objective.amici_model.setTimepoints(t)
for start in pypesto_result.optimize_result.list:
    par_dict = pypesto_problem.objective.par_arr_to_dct(np.power(10, start['x']))
    ignored_pars = set(par_dict.keys()).difference(pypesto_problem.objective.amici_model.getParameterIds())
    pypesto_problem.objective.amici_model.setParameterById({
        k: v
        for k, v in par_dict.items()
        if k not in ignored_pars
    })
    #pypesto_problem.objective.amici_model.setParameters(
    #    np.power(10, start['x']),
    #)
    #solver = model.getSolver()
    rdata = amici.runAmiciSimulation(
        pypesto_problem.objective.amici_model,
        pypesto_problem.objective.amici_solver,
    )
     
    #plotStateTrajectories(rdata, model=pypesto_problem.objective.amici_model)
    plotObservableTrajectories(
        rdata,
        model=pypesto_problem.objective.amici_model,
        observable_indices=observable_indices,
    )
    #plt.plot(t, [get_expected_x(t_i) for t_i in t])
    plt.show()
'''

#print(pypesto_problem.objective.amici_model.getObservableIds())
for start in pypesto_result.optimize_result.list[:1]:
    observable_indices=[pypesto_problem.objective.amici_model.getObservableIds().index('observable_Total_cumulative_cases')]
    par_dict = pypesto_problem.objective.par_arr_to_dct(np.power(10, start['x']))
    ignored_pars = set(par_dict.keys()).difference(pypesto_problem.objective.amici_model.getParameterIds())
    pypesto_problem.objective.amici_model.setParameterById({
        k: v
        for k, v in par_dict.items()
        if k not in ignored_pars
    })
    #pypesto_problem.objective.amici_model.setParameters(
    #    np.power(10, start['x']),
    #)
    #solver = model.getSolver()
    pypesto_problem.objective.amici_model.setTimepoints(np.linspace(0, 80, 801))
    rdata = amici.runAmiciSimulation(
        pypesto_problem.objective.amici_model,
        pypesto_problem.objective.amici_solver,
    )
     
    #plotStateTrajectories(rdata, model=pypesto_problem.objective.amici_model)
    fig, ax = plt.subplots(figsize=(20,10))
    plotObservableTrajectories(
        rdata,
        model=pypesto_problem.objective.amici_model,
        observable_indices=observable_indices,
        ax=ax,
    )
    #ax.set_yscale('log')
    #plt.plot(t, [get_expected_x(t_i) for t_i in t])
    plt.show()

pypesto_problem.objective.check_grad_multi_eps(x=petab_problem.x_nominal_free)



StarcoderdataJupyterScriptsDedupFiltered