armGraphs.py

# armGraphs.py -- Python code for the following purposes:
# (modified from testArm.py)
# *Read joint angles, muscle exc, muscle activation and force from .pnt files
# *Plot for each muscle: exc, act, force, muscle lengths (3 fibres+avg) 
# *Plot joint angles and trajectory over time
# Last update: 6/17/13 (salvadord)

import sys	#for exit
import struct
import time
import matplotlib
import matplotlib.pyplot as plt
import matplotlib.lines as lines
import matplotlib.text as text
import matplotlib.animation as animation
from mpl_toolkits.mplot3d import Axes3D
from pylab import figure, show
from numpy import * 
import csv
#import os


###############################
# Global constant parameters
###############################

#secLength = 6.01 # length of simulation in seconds (first 10 ms used to set start pos)
#msecInterval = 10.0 # interval at which packets are sent in miliseconds
#n = int(secLength*1000/msecInterval) # calculate number of samples
#toSamples = 1/msecInterval*1000 
numMuscles = 4 # number of muscles = shFlex (PECT), shExt (DELT), elFlex (BIC), elExt (TRIC)
#muscleNames =  ["Shoulder ext (post Deltoid+Infraspinatus)", "Shoulder flex (Pectoralis+ant Deltoid)", "Elbow ext (Triceps)", "Elbow flex (Biceps+Brachialis)"]
muscleNames =  ["Shoulder ext", "Shoulder flex", "Elbow ext", "Elbow flex"]
numMusBranches = 18 # number of muscle branches 
numJoints = 2 # number of joints (DOFs) 
armLen = [0.4634 - 0.173, 0.7169 - 0.4634] # elbow - shoulder from MSM;radioulnar - elbow from MSM;  useJointPos = 0 # use joint positions vs joint angles for 2d arm animation
showBranches = 0 # include muscle branches in graphs
verbose = 0 # whether to show output on screen 


###############################
# Function definition
###############################

	
# function to read the result .pnt file containing muscle excitation and force	
def readPntFiles(msmFolder, pntFile, secLength, msecInterval):
	
	n = int(secLength*1000/msecInterval) # calculate number of samples

	###################################
	# read joint position
	
	# open pnt file and read data
	fileName = msmFolder+"SUNY_arm_2DOFs_horizon_static_coordinate_status.pnt"
	try:
		#f = open(fileName, "r" )
		jointData = []	
		with open(fileName, "r") as f:
			#next(f)
			for line in f:
				print line
				jointData.append([float(i) for i in line.split()])
		jointData = array(jointData).transpose()
		jointData = jointData[:, len(jointData[0])-n:] # make number of rows equal to n (packets received) 
		
		# create arrays to store read data (joints include shoulder, elbow and wrist * 3 coords (xyz))
		jointPosSeq = zeros(((numJoints+1)*3, n)) 
		
		# assign activation and force to output arrays
		# file has format: time,ground_thorax_xyz,sternoclavicular_xyz,acromioclavicular_xyz,shoulder_xyz,elbow_xyz,radioulnar_xyz,radius_hand_xyz
		
		#jointPosSeq[0:6,:] = jointData[10:16,:]
		#jointPosSeq[6:9,:] = jointData[19:22,:]

		jointPosSeq[0,:] = jointData[10,:]
		jointPosSeq[1,:] = jointData[12,:]
		jointPosSeq[2,:] = jointData[13,:]
		jointPosSeq[3,:] = jointData[15,:]
		jointPosSeq[4,:] = jointData[19,:]
		jointPosSeq[5,:] = jointData[21,:]
	except:
		jointPosSeq=[]
		if verbose: 
			print "coordinate pnt file not available"

	########################################################
	# Read muscle activation and force
	
	# open pnt file and read data
	#fileName = msmFolder+"SUNY_arm_2DOFs_horizon_static_muscle_status.pnt"
	fileName = pntFile
	#f = open(fileName, "r" )
	muscleData = []	
	with open(fileName, "r") as f:
		#next(f)
		for line in f:
			muscleData.append([float(i) for i in line.split()])
	muscleData = array(muscleData).transpose() # transpose
	muscleData = muscleData[:, len(muscleData[0])-n:] # make number of rows equal to n (packets received) 
	
	# create arrays to store read data
	musExcSeq = zeros((numMusBranches, n))
	musActSeq = zeros((numMusBranches, n))
	musForcesSeq = zeros((numMusBranches, n))
	
	# assign activation and force to output arrays
	# file has format: time  DELT1_excitation  DELT1_activation  DELT1_force  DELT2_excitation  DELT2_activation  DELT2_force  ...
	musExcSeq = muscleData[1::3 , :]
	musActSeq = muscleData[2::3 , :]
	musForcesSeq = muscleData[3::3 , :]	
	
	return jointPosSeq, musExcSeq, musActSeq, musForcesSeq

# 	Plot for each muscle: exc, act, force, muscle lengths (3 fibres+avg), joint angles and trajectory over time
def plotGraphs(jointPosSeq, jointAnglesSeq, musLengthsSeq, musExcSeq, musActSeq, musForcesSeq, t1, t2, msecInterval, armAnimation, saveGraphs, saveName):
	# graph parameters
	toDegrees = 360/(2*pi)
	toSamples = 1/msecInterval*1000 
	legFont = 11
	linWidth = 3
	color1 = "red"
	color2 = "blue"
	color3 = "green"
	color4 = "purple"
	line1 = "-"
	line2 = ":"
	gridOn = 1
	
	# plot with 6 subplots for joint angles, trajectory and each of muscles
	fig1 = figure()
		
	# time variables
	T = arange(t1, t2, msecInterval/1000.0);
	t1Samples = t1*toSamples
	t2Samples = t2*toSamples
	
	if armAnimation:
		speedFactor = 1 # speed of animation
		
		#create figure and axes
		figAnim = figure()
		axAnim = figAnim.add_subplot(111)
		axAnim.grid(gridOn);
		# set axes size to fit arm
		axAnim.set_xlim(-armLen[0]-armLen[1]-0.1, armLen[0]+armLen[1]+0.1)
		axAnim.set_ylim(-armLen[0]-armLen[1]-0.1, armLen[0]+armLen[1]+0.1)
		
		###########################
		# 2D arm movement animation
		armImages=[]
		for t in arange(t1Samples,t2Samples):
			# Update the arm position based on jointPosSeq
			if useJointPos:
				# use x = z (pnt file) ; y = x (pnt file)
				shoulderPosx = jointPosSeq[1, t]
				shoulderPosy = jointPosSeq[0, t]
				elbowPosx = jointPosSeq[3, t]
				elbowPosy = jointPosSeq[2, t]
				wristPosx = jointPosSeq[5, t]
				wristPosy = jointPosSeq[4, t]
				
				# update jointAnglesSeq based on pos !!!
				
			# Update the arm position based on jointAnglesSeq	
			else:
				armAng = jointAnglesSeq[:,t]
				shoulderPosx = 0
				shoulderPosy = 0
				elbowPosx = armLen[0] * cos(armAng[0]) # end of elbow
				elbowPosy = armLen[0] * sin(armAng[0])
				wristPosx = elbowPosx + armLen[1] * cos(+armAng[0]+armAng[1]) # wrist=arm position
				wristPosy = elbowPosy + armLen[1] * sin(+armAng[0]+armAng[1])
			
			# create
			armLine1 = lines.Line2D([0, elbowPosx-shoulderPosx], [0, elbowPosy-shoulderPosy], color=color1, linestyle=line1, linewidth=linWidth)
			armLine2 = lines.Line2D([elbowPosx-shoulderPosx, wristPosx-shoulderPosx], [elbowPosy-shoulderPosy, wristPosy-shoulderPosy], color=color2, linestyle=line1, linewidth=linWidth)
			axAnim.add_line(armLine1)
			axAnim.add_line(armLine2)
			#label = plt.legend(armLine1, str(t/toSamples) )
			label = text.Text(x=0, y=0.5, text="time = "+str(t/toSamples), weight='bold' )
			axAnim.add_artist(label)
			armImages.append([armLine1, armLine2, label])
		
		# add blank frames
		blankFrames = int(1*toSamples)
		for t in range(blankFrames):
			# Update the arm position
			armLine1 = lines.Line2D([0, 0], [0, 0], color=color1, linestyle=line1, linewidth=linWidth)
			armLine2 = lines.Line2D([0,0 ], [0,0], color=color2, linestyle=line1, linewidth=linWidth)
			axAnim.add_line(armLine1)
			axAnim.add_line(armLine2)			
			armImages.append([armLine1, armLine2])
		
		# generate animation 
		armAnim = animation.ArtistAnimation(figAnim, armImages, interval=msecInterval/speedFactor, repeat_delay=500, blit=True)
	
	###########################
	# Plot joint angles vs time
	ax = fig1.add_subplot(321)
	#T = arange(0, secLength, msecInterval/1000.0);
	
	T=T[:len(jointAnglesSeq[0,t1Samples:t2Samples])]
	ax.plot(T,jointAnglesSeq[0,t1Samples:t2Samples]*toDegrees,color=color1,linestyle=line1, linewidth=linWidth, label="shoulder")
	ax.plot(T,jointAnglesSeq[1,t1Samples:t2Samples]*toDegrees,color=color2,linestyle=line1, linewidth=linWidth, label="elbow")
      
	ax.set_ylabel('angle (deg)', fontsize = legFont)
	ax.set_xlabel('time (sec)', fontsize = legFont)
	ax.set_title('Joint angles')
	ax.legend(loc='upper center', bbox_to_anchor=(1.0, 1.0),  borderaxespad=0., prop={'size':legFont})
	#ax.set_xlim([t1, t2])
    #ax.set_ylim(bmmYlims_sh)
	ax.grid(gridOn)
	
	############################
	# Plot x-y pos vs time
	ax = fig1.add_subplot(322)
	
	# calculate x and y trajectories based on angles
	if useJointPos:
		xTraj = jointPosSeq[4, t1Samples:t2Samples]
		yTraj = jointPosSeq[5, t1Samples:t2Samples]
	else:
		xTraj = armLen[0]*cos(jointAnglesSeq[0,t1Samples:t2Samples])+armLen[1]*cos(jointAnglesSeq[1,t1Samples:t2Samples])
		yTraj = armLen[0]*sin(jointAnglesSeq[0,t1Samples:t2Samples])+armLen[0]*sin(jointAnglesSeq[1,t1Samples:t2Samples])
	
	#ax.plot(xTraj, yTraj,color=color2,linestyle=line1, linewidth=linWidth)
	ax.plot(T, xTraj,color=color1,linestyle=line1, linewidth=linWidth, label="x")
	ax.plot(T, yTraj,color=color2,linestyle=line1, linewidth=linWidth, label="y")
	
	ax.set_ylabel('position (m)', fontsize = legFont)
	#ax.set_xlabel('x position (m)', fontsize = legFont)
	ax.set_xlabel('time (sec)', fontsize = legFont)
	ax.set_title('X-Y trajectory')
	ax.legend(loc='upper center', bbox_to_anchor=(1.0, 1.0),  borderaxespad=0., prop={'size':legFont})
	#ax.set_xlim([t1, t2])
    #ax.set_ylim(bmmYlims_sh)
	ax.grid(gridOn)
	
	############################
	# Plot excitation, activation and force for each muscle
	
	# calculate normalized force (activation and length already normalized)
	#musActSeqNorm = musActSeq/musActSeq[:,t1Samples:t2Samples].max()
	musForcesSeqNorm = musForcesSeq/musForcesSeq[:,t1Samples:t2Samples].max()
	#musLengthsSeqNorm = musLengthsSeq/musLengthsSeq[:,t1Samples:t2Samples].max()
	
	# Note arrangement of muscle branches in data arrays:
	#DELT1(0)  DELT2(1) DELT3(2) Infraspinatus(3) Latissimus_dorsi_1(4) Latissimus_dorsi_2(5) Latissimus_dorsi_3(6) 
	#Teres_minor(7) PECM1(8) PECM2(9) PECM3(10) Coracobrachialis(11) TRIlong(12) TRIlat(13) TRImed(14) BIClong(15) BICshort(16) BRA(17)  
	# is different from muscle groups:
	# Sh ext = DELT3, Infraspinatus, Latissimus_dorsi_1-3, Teres_minor
	# Sh flex = PECM1, PECM2, PECM3, DELT1, Coracobrachialis
	# El ext = TRIlong, TRIlat, TRImed
	# El flex = BIClong, BICshort, BRA
	shext=[2,3,4,5,6,7]
	shflex=[0,8,9,10,11]
	elext=[12,13,14]
	elflex=[15,16,17]
	musList=[shext, shflex,elext,elflex]
		
	for iMus in range(numMuscles):
		ax = fig1.add_subplot(3,2,iMus+3)
		# set number of muscle branches - assume each node has 3 branches (treat the Brachialis as a branch of Biceps=elbow flexor)
		#iBranches = 3
		
		### Excitation and Activation ####
		# equivalent for all branches of same muscle group
		offset = 2  # use offset 3 because only DELT3 is used (order of muscle branches doesn't correspond muscle groups!)
		ax.plot(T, musExcSeq[musList[iMus][offset],t1Samples:t2Samples],color=color1,linestyle=line1, linewidth=linWidth, label="excitation")
		ax.plot(T, musActSeq[musList[iMus][offset],t1Samples:t2Samples],color=color2,linestyle=line1, linewidth=linWidth, label="activation")
		
		# for show branches plot individual branches and mean value for force and length
		if showBranches:
			### Force and Length ###
			for iBranch in range(len(musList[iMus])):
				ax.plot(T, musForcesSeqNorm[musList[iMus][iBranch],t1Samples:t2Samples],color=color3,linestyle=line2, linewidth=linWidth-1)
				ax.plot(T, musLengthsSeq[musList[iMus][iBranch],t1Samples:t2Samples],color=color4,linestyle=line2, linewidth=linWidth-1)
			ax.plot(T, musForcesSeqNorm[musList[iMus],t1Samples:t2Samples].mean(axis=0),color=color3,linestyle=line1, linewidth=linWidth, label="force (mean)")
			ax.plot(T, musLengthsSeq[musList[iMus],t1Samples:t2Samples].mean(axis=0),color=color4,linestyle=line1, linewidth=linWidth, label="length (mean)")
		# for NOT show branches show mean value for force and single value for length
		else:
			### Force ###
			# For shoulder extensor group show only posterior Deltoid, branch 3 (DELT3) or Infraspinatus (INFSP) 
			# branch 2 (DELT2 = lateral deltoid) also available but currently not included in shoulder extensor group
			if iMus == 0:
				offset = 2 # DELT3
				#offset = 12 # INFSP 
				ax.plot(T, musForcesSeqNorm[musList[iMus][offset],t1Samples:t2Samples],color=color3,linestyle=line1, linewidth=linWidth, label="force")
			# For rest of muscles use mean value of all branches
			else:
				offset=0
				ax.plot(T, musForcesSeqNorm[musList[iMus],t1Samples:t2Samples].mean(axis=0),color=color3,linestyle=line1, linewidth=linWidth, label="force")
			
			### Length ####
			# show length only of one muscle indicated by the index 'offset' - DELT3, PECM1, BIClong, TRIlong
			maxLength = 0.20 
			ax.plot(T, musLengthsSeq[musList[iMus][offset],t1Samples:t2Samples]/maxLength,color=color4,linestyle=line1, linewidth=linWidth, label="length")
	
		# show branche label
		if (showBranches):
			ax.plot(-1, -1,color=color3,linestyle=line2, linewidth=linWidth, label="force (branches)")
			ax.plot(-1, -1,color=color4,linestyle=line2, linewidth=linWidth, label="length (branches)")
		
		# axis properties
		ax.set_ylabel('normalized value', fontsize = legFont)
		ax.set_ylim([0,1])
		ax.set_xlim([t1,t2])
		ax.set_xlabel('time (sec)', fontsize = legFont)
		ax.set_title(muscleNames[iMus])
		if iMus==3:
			ax.legend(loc='upper center', bbox_to_anchor=(-0.2, 1.8),  borderaxespad=0., prop={'size':legFont})
		ax.grid(gridOn)
	
	# show graphs
	fig1.tight_layout()
	show()
	
	# save graphs using startPos and pattern in filename
	if saveGraphs:
		saveFolder = 'gif/'
		fig1.savefig(saveFolder+saveName, bbox_inches=0)

		#if armAnimation:
			#armAnim.save('test.mp4')
			#armAnim.save(saveFolder+saveName+'.mp4',writer = writer)
	

# run single test (udp transfer, read files, plot graphs)	
def readAndPlot(jointAnglesSeq, musLengthsSeq, msmFolder, armAnimation, saveGraphs, saveName, timeRange, msecInterval):
	# Sim parameters
	#armAnimation = 1 #  show 2D arm animation
	#saveGraphs = 1 # save graph and animation 
	
	# define time interval to display
	#timeInterval = [0.1, 30]
	
	# Send muscle excitations to MSM and receive joint angles and muscle lengths
	#jointAnglesSeq, musLengthsSeq = sendAndReceiveMsmData(initJointAngles, musExcSeq, readSimFromFile)
	
	# Read data from .pnt files
	jointPosSeq,musExcSeq, musActSeq, musForcesSeq = readPntFiles(msmFolder, timeRange[1], msecInterval)
	
	# Plot results (last 2 arguments = initial and end times in seconds)
	plotGraphs(jointPosSeq, jointAnglesSeq, musLengthsSeq, musExcSeq, musActSeq, musForcesSeq, timeRange[0], timeRange[1], msecInterval, armAnimation, saveGraphs, saveName)

##############################
# Main script
##############################
'''
jointAnglesSeq = zeros((numJoints, n))
musLengthsSeq = zeros((numMusBranches, n))
armAnimation = 1 #  show 2D arm animation
saveGraphs = 1 # save graph and animation 
timeInterval = [0.1, 30]
saveName='temp'
msmFolder = "/home/salvadord/Documents/ISB/Models_linux/msarm/source/test/"

readAndPlot(jointAnglesSeq, musLengthsSeq, msmFolder, armAnimation, saveGraphs, saveName, timeInterval)
'''