Zimm–Bragg Model calculator

Zimm–Bragg.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
'''Zimm–Bragg Model calculator'''
import itertools
LENGTH =9

def weight(weight_sigma = 0, weight_s = 0):
	if weight_sigma == 0 and weight_s == 0:
		return '1'
	elif weight_sigma == 0:
		return 's'+str(weight_s)
	elif weight_s == 0:
		return 'σ'+str(weight_sigma)
	else:
		return 'σ'+str(weight_sigma)+'s'+str(weight_s)

sequences = itertools.product('CH', repeat = LENGTH)
total_weight = {}
Hbond_dict = {}
print 'Seq.\tIgnored?\tWeight\t#Hbond'
print '----\t----\t----\t----'
for sequence in sequences:
	sequence_str = ''.join(sequence)
	# if ignore?
	ignored = False
	if (sequence_str.startswith('H') and sequence_str.find('CH')>-1) or (sequence_str.count('CH')>1):
		ignored = True
	# stat weight:
	weight_sigma = 0
	weight_s = 0
	previous_residue = ''
	for residue in sequence:
		if residue == 'H':
			if previous_residue!='H':
				weight_sigma += 1
				weight_s += 1
			else: # previous_residue=='H'
				weight_s += 1
		previous_residue = residue
	this_weight = weight(weight_sigma, weight_s)
	# ignored_str
	if ignored:
		ignored_str = 'y'
	else:
		ignored_str = 'n'
	# main
	if not ignored:
		# Hbond_amount:
		Hbond_amount = sequence_str.count('H') - 2
		if Hbond_amount<0:
			Hbond_amount = 0
		Hbond_dict[this_weight] = Hbond_amount
		# register this weight contribution:
		if total_weight.has_key(this_weight):
			total_weight[this_weight] += 1
		else:
			total_weight[this_weight] = 1
		# output:
		print sequence_str+'\t'+ignored_str+'\t'+this_weight+'\t'+str(Hbond_amount)

print '----------------'

averg_HbondAmount_str = ''
total_weight_str = ''
for k,v in total_weight.items():
	if v == 1:
		v_str = ''
	else:
		v_str = str(v)
	total_weight_str += v_str+' '+k+' + '
	if Hbond_dict[k]>0:
		averg_HbondAmount_str += str(Hbond_dict[k])+'*'+v_str+' '+k+' + '
total_weight_str = total_weight_str[:-3] # to remove last " + "
averg_HbondAmount_str = averg_HbondAmount_str[:-3] # to remove last " + "
print '  z =',total_weight_str
print '<n> =', averg_HbondAmount_str
print '\n     ', averg_HbondAmount_str
print '  θ = ---------------------------------'
print '                    ', LENGTH-2