source: branches/numpy_misc/tools/pytools/numtools.py @ 6818

"""Functions for numerical computations
"""

epsilon = 1.0e-15





def sum(x):
  """
    Attempt to sum up elements in x
  """

  import types
  import numpy as num
  
  if isinstance(x, num.ndarray):
    return num.sum(x)
  elif isinstance(x, (list, tuple)):
    s = x[0]
    for e in x[1:]:
      s += e
    return s  

def mvmul(A, x):
  """Multiply matrix A onto vector x
  """

  import numpy as num

  x = num.reshape(x, (A.shape[1], 1)) #Make x a column vector             

  return num.dot(A, x)


def all_equal(vec):

  equal = 1
  v0 = vec[0]
  for v in vec:
    if v != v0:
      equal = 0   

  return equal


def meshgrid(N):
  """Make meshgrid (a' la Matlab)  
  """
  
  import numpy as num

  d = len(N)

  X = [] 
  for s in range(d):
    local_shape = num.ones(d)
    local_shape[s] = N[s]
    a = num.arange(N[s])
    if N[s] > 1:
      a = a.astype(num.float)/(N[s]-1)
    else:
      a = a.astype(num.float)
    
    # Put ones in all other dimensions
    #
    e = [] 
    for t in range(d):
      if s == t:  
        e.append(a)
      else:
        e.append(num.ones(N[t]))
    
    # Take kronecker product of all dimensions
    #
    x = 1  
    for t in range(d):
      x=num.multiply.outer(x,e[t])

    #print x, x.shape
    X.append(x)
    
  return X    
    

def expand(x, mask):
  """Expand vector x into into vector of length equal to vector
     mask such that elements of x are placed where mask is one.
     Number of ones in mask must equal len(x)."""

  import numpy as num

  assert isinstance(x, num.ndarray)
  assert isinstance(mask, num.ndarray)
  #FIXME: Assert that mask contains only ones and zeros
  assert len(x) == num.sum(mask), 'Number of ones in mask must equal length of x'

  d = len(mask)
  y = num.zeros(d, x.dtype)

  i = 0
  for s in range(d):
    if mask[s]:
      y[s] = x[i] 
      i += 1

  return y    


def pwr2trunc(N, e = None, dir = 0): 
    """N1 = pwr2trunc(N, e, dir)

    If e is None, let e be the largest integer such that N > 2**e.
    
    If dir = 0 (default)
      Compute the nearest number smaller than N divisible by 2^e
    if dir == 1
      Compute the nearest number greater than N divisible by 2^e

     
    """  

    import math

    if e is None:
        e = int(math.log(N)/math.log(2))  # Maximal exponent

    k = N % 2**e
    N1 = N - k
    if dir == 1:
        N1  = 2*N1

    return N1, e