##########################################################################
# CalcUtils.py:  Python source file for the Simcal1 module.              #
# Copyright (c) 2004 Daniel Dvorkin.                                     #
# Distributed under the Artistic License, WITHOUT "optional clause 8"    #
# (see http://www.opensource.org/licenses/artistic-license.php).         #
# All rights not granted in this license are reserved by the author.     #
##########################################################################

"""Defines some common mathematical and statistical functions."""

# MASTER IMPORT STATEMENT

from Simcal1 import *


# FUNCTION DEFINITIONS

def prange(x):
	"""Positive range:  returns xrange of 1 to x inclusive."""
	
	return xrange(1, x+1)

def sample_std(x):
	"""Biased estimator for the standard deviation of x -- i.e., division
	is by n, the number of observations, rather than by n-1.  Differs from
	MLab.std() in two other important ways.	 First, data are presented in 
	row-variable, column-observation fashion.  Second, no type conversion 
	is performed -- x must already be a Numeric array."""
	
	n = float(x.shape[1])
	mx = mean(x, 1)
	mx.shape = mx.shape[:1] + (1,) + mx.shape[1:]
	
	x = x - mx
	
	return squeeze(sqrt(add.reduce(x * x, 1) / n))
	
def sample_cov(x, y):
	"""Biased estimator for the covariance matrix of x and y -- i.e.,
	normalization is by n, the number of observations, rather than by n-1.
	Differs from MLab.cov() in some other important ways:  Both x and y
	must be passed as arguments.  Data are presented in row-variable,
	column-observation fashion.	 No type-checking or shape-checking is
	performed.	All arguments must be real."""
	
	n = float(x.shape[1])
	mx = mean(x, 1)
	mx.shape = mx.shape[:1] + (1,) + mx.shape[1:]
	my = mean(y, 1)
	my.shape = my.shape[:1] + (1,) + my.shape[1:]
	
	x = x - mx
	y = y - my
	
	return squeeze(dot(x, transpose(y)) / n)
	
def corr_coef(x, y):
	"""Calculates the correlation coefficient(s) of x and y, but more 
	robust than MLab.corrcoef() in that the lengths of the x and y
	matrices need not be equal.	 Data are presented as row-variable,
	column-observation."""
	
	sx = sample_std(x)
	sy = sample_std(y)
	sps = outerproduct(sx, sy)
	
	cvs = sample_cov(x, y)
	
	return squeeze(cvs / sps)
	
def uc_corr_coef(x, y):
	"""Calculates the uncentered correlation coefficient(s) of x and y,
	defined as the cosine of the angles between the vector(s) making up
	x and y.  The lengths of the x and y matrices need not be equal, but
	they must have the same number of observations per variable.  Data are
	presented as row-variable, column-observation."""
	
	x_mags = sqrt(sum(x**2, 1))
	y_mags = sqrt(sum(y**2, 1))
	
	return squeeze(dot(x, transpose(y)) / outerproduct(x_mags, y_mags))
	
def max_of(x, y):
	"""Given either two scalars, two Numeric arrays of equal shape, or
	a Numeric array and a scalar, return the elementwise maxima.  E.g., if
	a = array([1, 2, 3]) and b = 2, then max_of(a, b) will return
	array([2, 2, 3])."""
	
	# find max_of(x-y, 0) because this is easier
	x = x - y
	x = x * (x > 0)
	
	# now we add y back in
	return x + y
		
def max_zero(x):
	"""Special case of max_of.	Given either a scalar or a Numeric array
	x, return the elementwise maxima of x and 0."""
	
	return x * (x > 0)

def array_split(arr, elems):
	"""Takes a Numeric array "arr" and returns the array split into a list
	of "elems" chunks.  E.g.,
		
		a = array([1,2,3,4,5,6])
		b = array_split(a, 3)
	
	leads to:
	
		b = array([[1,2],[3,4],[5,6]])"""
	
	length = arr.shape[0]
	elsize = int(ceil(float(length) / float(elems)))
	retn = [None] * elems
	
	for e in xrange(0, elems):
		if e < elems-1:
			retn[e] = arr[e*elsize : (e+1)*elsize]
		else:
			retn[e] = arr[e*elsize :]
	
	return retn