source: tools/pytools/numtools.py @ 2017

"""Functions for numerical computations
"""

epsilon = 1.0e-15





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

  import types, Numeric
  
  if type(x) == Numeric.ArrayType:
    return Numeric.sum(x)
  elif type(x) in [types.ListType, types.TupleType]:
    s = x[0]
    for e in x[1:]:
      s += e
    return s  

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

  from Numeric import dot, reshape

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

  return 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)  
  """
  
  from Numeric import ones, arange, Float, multiply

  d = len(N)

  X = [] 
  for s in range(d):
    local_shape = ones(d)
    local_shape[s] = N[s]
    a = arange(N[s])
    if N[s] > 1:
      a = a.astype(Float)/(N[s]-1)
    else:
      a = a.astype(Float)
    
    # Put ones in all other dimensions
    #
    e = [] 
    for t in range(d):
      if s == t:  
        e.append(a)
      else:
        e.append(ones(N[t]))
    
    # Take kronecker product of all dimensions
    #
    x = 1  
    for t in range(d):
      x=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)."""

  from Numeric import sum, ArrayType, zeros

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

  d = len(mask)
  y = zeros(d, x.typecode())

  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