shear_file.py

'''
Created on 24/04/2014

@author: MMPE
'''
from __future__ import unicode_literals
from __future__ import print_function
from __future__ import division
from __future__ import absolute_import
from builtins import range
from io import open
from future import standard_library
standard_library.install_aliases()

import numpy as np
import os

def save(filename, x_coordinates, z_coordinates, u=None, v=None, w=None):
    """
    Parameters
    ----------
    filename : str
        filename
    x_coordinates : array_like
        lateral coordinates
    z_coordinates : array_like
        vertical coordinates
    u : array_like, optional
        shear_u component, normalized with U_mean\n
        shape must be (#z_coordinates, #x_coordinates) or (#z_coordinates,)
    v : array_like, optional
        shear_v component, normalized with U_mean\n
        shape must be (#z_coordinates, #x_coordinates) or (#z_coordinates,)
    w : array_like, optional
        shear_w component, normalized with U_mean\n
        shape must be (#z_coordinates, #x_coordinates) or (#z_coordinates,)
    """

    shape = (len(z_coordinates), len(x_coordinates))
    vuw = [v, u, w]
    for i in range(3):
        if vuw[i] is None:
            if i == 1:
                vuw[i] = np.ones((shape))
            else:
                vuw[i] = np.zeros((shape))
        else:
            vuw[i] = np.array(vuw[i])
            if len(vuw[i].shape) == 1 and vuw[i].shape[0] == shape[0]:
                vuw[i] = np.repeat(np.atleast_2d(vuw[i]).T, shape[1], 1)

            assert vuw[i].shape == shape, (i, vuw[i].shape, shape)

    # exist_ok does not exist in Python27
    if not os.path.exists(os.path.dirname(filename)):
        os.makedirs(os.path.dirname(filename))#, exist_ok=True)
    with open(filename, 'w', encoding='utf-8') as fid:
        fid.write(" # autogenerated shear file\n")
        fid.write("  %d %d\n" % (len(x_coordinates), len(z_coordinates)))
        for i, l in enumerate(['v', 'u', 'w']):
            fid.write(" # shear %s component\n  " % l)
            fid.write("\n  ".join([" ".join(["%.10f" % v for v in r ]) for r in vuw[i]]))
            fid.write("\n")
        for yz, coor in (['v', x_coordinates], ['w', z_coordinates]):
            fid.write(" # %s coordinates\n  " % yz)
            fid.write("\n  ".join("%.10f" % v for v in coor))
            fid.write("\n")


if __name__ == "__main__":
    save("test.dat", [-55, 55], [30, 100, 160] , u=np.array([[0.7, 1, 1.3], [0.7, 1, 1.3]]).T)