Plotting a tri-trophic system in PyX

tritrophicPlot.py

from pyx import *
import numpy as np
import scipy as sp

def null_bernoulli(top,bot,nint):
	tint = top*bot
	zeroes = tint-nint
	## We only return a matrix with no non-interacting species
	hasNull = True
	trials = 0
	while hasNull:
		trials += 1
		mat = np.concatenate((np.ones(nint),np.zeros(zeroes)),1)
		np.random.shuffle(mat)
		mat = mat.reshape((top,bot))
		NI0 = 0
		NI1 = 0
		for ms0 in mat.sum(axis=0):
			if ms0 == 0:
				NI0 = NI0 + 1
		for ms1 in mat.sum(axis=1):
			if ms1 == 0:
				NI1 = NI1 + 1
		if (NI0 == 0)&(NI1 == 0):
			hasNull = False
	return mat


Nr = 2
Nc = 4
Np = 3

A = null_bernoulli(Np,Nc,np.floor(Np*Nc*0.8))
C = null_bernoulli(Nc,Nr,np.floor(Nr*Nc*0.5))

def rank(V):
	# Returns the rank of a vector
	# with no ties
	rn = np.zeros(len(V),dtype=np.int32)
	crnk = 0
	while crnk < len(V):
		for j in range(0,len(V)):
			cMax = max(V)
			if V[j] == cMax:
				rn[j] = crnk
				crnk += 1
				V[j] = min(V)-1
				break
	return rn


def findXpos(n,rad):
	if n == 1:
		pos = [0]
	if n == 2:
		pos = [-2*rad,2*rad]
	if n > 2:
		pos = np.linspace(0-n*2*rad,0+n*2*rad,n)
	return pos


def shortenLine(x1,y1,x2,y2,d):
	a = (y2-y1)/(x2-x1)
	b = y1-a*x2
	od = np.sqrt(np.power((x2-x1),2)+np.power((y2-y1),2))
	R = (od-d)/float(od)
	x1b = R*(x2-x1)+x1
	x2b = R*(x1-x2)+x2
	y1b = R*(y2-y1)+y1
	y2b = R*(y1-y2)+y2
	return [x1b,y1b,x2b,y2b]

def plotTr(A,C,file='tritrophic',suppl=[],sprad=0.5,elevfac=5,shfac=1.2):
	# Get the cumulative degrees, ranks, and number of species
	Dm = np.sum(C,axis=1)+np.sum(A,axis=0)
	Dl = np.sum(C,axis=0)
	Dt = np.sum(A,axis=1)
	Rm = rank(Dm)
	Rl = rank(Dl)
	Rt = rank(Dt)
	Nm = len(Dm)
	Nl = len(Dl)
	Nt = len(Dt)
	## Begin plot
	c = canvas.canvas()
	# If there are any supplementary arguments
	if len(suppl)>0:
		text.set(mode="latex")
		for suarg in suppl:
			text.preamble(suarg)
	# End of supplementary arguments
	## General arguments
	Yl = 0
	Ym = Yl + sprad * elevfac
	Yt = Ym + sprad * elevfac
	Xl = findXpos(Nl,sprad)
	Xm = findXpos(Nm,sprad)
	Xt = findXpos(Nt,sprad)
	## Plot arrows R -> C
	for ls in range(len(Dl)):
		for ms in range(len(Dm)):
			if C[ms,ls] == 1:
				co = shortenLine(Xl[Rl[ls]], Yl, Xm[Rm[ms]], Ym, sprad*shfac)
				c.stroke(path.line(co[0],co[1],co[2],co[3]),[deco.earrow(),deco.barrow()])
	## Plot arrows C -> P
	for ms in range(len(Dm)):
		for ts in range(len(Dt)):
			if A[ts,ms] == 1:
				co = shortenLine(Xm[Rm[ms]], Ym, Xt[Rt[ts]], Yt, sprad*shfac)
				c.stroke(path.line(co[0],co[1],co[2],co[3]),[deco.earrow(),deco.barrow()])
	## Plot lower trophic level
	for xl in Xl:
		c.fill(path.circle(xl, Yl, sprad),[deco.filled([color.cmyk.ForestGreen])])
	## Plot middle trophic level
	for xl in Xm:
		c.fill(path.circle(xl, Ym, sprad),[deco.filled([color.cmyk.NavyBlue])])
	## Plot upper trophic level
	for xl in Xt:
		c.fill(path.circle(xl, Yt, sprad),[deco.filled([color.cmyk.BurntOrange])])
	c.writePDFfile(file)
	## End plot
	return 0


plotTr(A,C,suppl=['\usepackage[T1]{fontenc}','\usepackage{fourier}','\usepackage[lf]{venturis}'])