tidal_tails/utils.py

"""Short utility functions for use elsewhere"""

import numpy as np

def update_config(default, update):
    """Recursively updates a Python (configuration) dictionary"""
    for key in update: 
        if isinstance(default[key], dict): #Recurse if it is a dictionary
            update_config(default[key], update[key])
        else: default[key] = update[key] #Simply copy the value otherwise

def random_unit_vectors(size, mode='3D'):
    """Draws an array of size size corresponding to isotropically distributed
    random unit vectors (number of vectors, dimension). If mode is 2D, the
    vectors are instead distributed isotropically on the xy plane."""
    direction = np.random.normal(size=(size, 3))
    if mode == '2D': direction[:, -1] = 0
    return direction / np.linalg.norm(direction, axis=-1, keepdims=True)

def cascade_round(arr):
    """Rounds the floats in an array to integers while preserving their total
    sum (as closely as possible). Follows the cascade rounding algorithm."""
    runningSurplus = 0
    for i in range(len(arr)):
        f = arr[i] + runningSurplus
        arr[i] = int(np.round(f))
        runningSurplus = f - arr[i]
    return arr