#!/usr/bin/env python
# coding: utf-8

# # Test tube analysis example
# 
# Analyze a test tube containing 2 strand species that interact to form all complex species of up to 4 strands. 
# 
# Strand a: 5e-6 M    
# Strand b: 5e-6 M
# 
# Strands a and b are intended to form a bipedal walker complex '(a+b)' with a duplex torso and single-stranded legs (Shin and Pierce, J Am Chem Soc, 2004). 
# 
# Material: DNA   
# Temperature: 23 C      
# 
# Calculate the partition function, equilibrium pair probability matrix, MFE proxy structure(s), and equilibrium concentration for each complex species in the tube, as well as the ensemble pair fractions for the tube ensemble. 

# In[1]:


# Import NUPACK Python module
from nupack import *


# In[2]:


# Define physical model
my_model = Model(material='dna', celsius=23)

# Define strand species
a = Strand('GGCTGGTTTCTGCTCTCTAGTTCGCGAGGTGCAATCTCCTATC', name='a')
b = Strand('GTCTGGGATGCTGGATACTGAACCTAGAGAGCAGAAACCAGCC', name='b')

# Define tube ensemble containing strands at specified concentrations 
# interacting to form all complexes up to 4 strands
t1 = Tube(strands={a:5e-6, b:5e-6}, complexes=SetSpec(max_size=4), name='Tube t1')

# Analyze the tube ensemble
# Calculate pfunc (default), pairs, mfe, concentration (default) for each complex 
# Since pairs is specified, calculate ensemble pair fractions for the tube ensemble
tube_result = tube_analysis(tubes=[t1], compute=['pairs', 'mfe'], model=my_model)
tube_result


# In[3]:


# MFE proxy structure for bipedal walker complex '(a+b)'
walker_result = tube_result['(a+b)']
print('\nMFE proxy structure(s) for walker (a+b):')
for i, s in enumerate(walker_result.mfe):
    print('    %2d: %s (%.2f kcal/mol)' % (i, s.structure, s.energy))


# In[4]:


# Plot the equilibrium pair probability matrix for bipedal walker complex '(a+b)'
import matplotlib.pyplot as plt
plt.imshow(walker_result.pairs.to_array())
plt.xlabel('Base index')
plt.ylabel('Base index')
plt.title('Pair probabilities for complex (a+b)')
plt.colorbar()
plt.clim(0, 1)


# In[5]:


# Plot the ensemble pair fractions for the tube 
plt.imshow(tube_result[t1].ensemble_pair_fractions.to_array())
plt.xlabel('Base index')
plt.ylabel('Base index')
plt.title('Ensemble pair fractions for ' + ', '.join(s.name for s in t1.strands))
plt.colorbar()
plt.clim(0, 1)