import numpy
from pysb.bng import generate_equations
from pysb.simulator import ScipyOdeSimulator
from pysb import Observable, ReactionPattern

from util import pkldump, pklload
from process_data import *

# ### User defined parameters
time_lim_hrs = 10
time_steps = 100
# ##########################


def run_simulation_condition(model, ts, inhib_monomers):
    if inhib_monomers:
        print('Setting inhibition condition for %s' %
              ', '.join(inhib_monomers))
    old_init = {}
    for monomer_name in inhib_monomers:
        init_name = monome_name + '_0'
        old_init[monomer_name] = model.parameters[init_name].value
        model.parameters[init_name].value = 0

    print('Instantiating simulator')
    sim = ScipyOdeSimulator(model, ts)

    print('Running simulation')
    res = sim.run()

    print('Reverting change to model initial conditions')
    for monomer_name in inhib_monomers:
        init_name = monome_name + '_0'
        model.parameters[init_name].value = old_init[monomer_name]
    return res


def set_model_observables(model):
    s6 = model.monomers['RPS6']
    site_patterns = [
        {'S235': 'p', 'S236': 'u'}
        {'S235': 'u', 'S236': 'p'}
        {'S235': 'p', 'S236': 'p'}
        ]
    pattern_terms = [s6(**pat) for pat in site_patterns]
    obs_pattern = ReactionPattern(pattern_terms)
    obs = Observable('pS6', obs_pattern)
    model.add_component(obs)


if __name__ == '__main__':
    time_lim_seconds = time_lim_hrs * 3600
    ts = numpy.linspace(0, time_lim_seconds, time_steps + 1)

    for cell_line  in ('C32', 'LOXIMVI', 'MMACSF', 'MZ7MEL', 'RVH421'):
        print('Loading model for %s' % cell_line)
        model = pklload('pysb_model_%s' % cell_line)
        print('Generating model equations for %s' % cell_line)
        generate_equations(model)
        res_nodrug = run_simulation_condition(model, ts, [])
        res_drug = run_simulation_condition(model, ts, ['ARAF', 'BRAF', 'RAF1'])
        break


@@ -1,12 +1,24 @@
+import copy%0A
 import numpy
@@ -147,16 +147,32 @@
 nPattern
+, ComplexPattern
 %0A%0Afrom u
@@ -352,25 +352,19 @@
 ndition(
-model, ts
+sim
 , inhib_
@@ -504,17 +504,103 @@
 old_
-init = %7B%7D
+params = %7Bp.name: p.value for p in sim.model.parameters%7D%0A    sim_params = copy.copy(old_params)
 %0A   
@@ -654,32 +654,33 @@
 it_name = monome
+r
 _name + '_0'%0A   
@@ -688,48 +688,17 @@
     
-old_init%5Bmonomer_name%5D = model.
+sim_
 param
-eter
 s%5Bin
@@ -709,56 +709,8 @@
 ame%5D
-.value%0A        model.parameters%5Binit_name%5D.value
  = 0
@@ -726,18 +726,12 @@
 nt('
-Instantiat
+Runn
 ing 
@@ -741,103 +741,55 @@
 ulat
-or
+ion
 ')%0A    
-sim = ScipyOdeSimulator(model, ts)%0A%0A    print('Running simulation')%0A    res = sim.run(
+res = sim.run(param_values=sim_params
 )%0A%0A 
@@ -911,16 +911,17 @@
 = monome
+r
 _name + 
@@ -977,20 +977,19 @@
 old_
-init%5Bmonomer
+params%5Binit
 _nam
@@ -1042,16 +1042,55 @@
 model):%0A
+    print('Setting model observables')%0A
     s6 =
@@ -1134,25 +1134,16 @@
 erns = %5B
-%0A        
 %7B'S235':
@@ -1160,17 +1160,31 @@
 6': 'u'%7D
+,
 %0A
+             
         
@@ -1201,33 +1201,47 @@
 u', 'S236': 'p'%7D
+,
 %0A
+             
         %7B'S235':
@@ -1258,25 +1258,16 @@
 6': 'p'%7D
-%0A        
 %5D%0A    pa
@@ -1285,17 +1285,62 @@
  = %5B
-s6(**pat)
+ComplexPattern(%5Bs6(**pat)%5D, None)%0A                    
  for
@@ -1450,16 +1450,31 @@
 _pattern
+, _export=False
 )%0A    mo
@@ -1797,32 +1797,71 @@
 s' %25 cell_line)%0A
+%0A        set_model_observables(model)%0A%0A
         print('G
@@ -1885,36 +1885,17 @@
 quations
- for %25s' %25 cell_line
+'
 )%0A      
@@ -1922,24 +1922,25 @@
 (model)%0A
+%0A
         
 res_nodr
@@ -1935,60 +1935,115 @@
     
-res_nodrug = run_simulation_condition(model, ts, %5B%5D)
+print('Instantiating simulator')%0A        sim = ScipyOdeSimulator(model, ts, use_theano=True, verbose=True)%0A
 %0A   
@@ -2083,25 +2083,19 @@
 ndition(
-model, ts
+sim
 , %5B'ARAF
@@ -2114,16 +2114,71 @@
 RAF1'%5D)%0A
+        res_nodrug = run_simulation_condition(sim, %5B%5D)%0A
         

Improve simulation script for Task 1