diff --git a/wetb/hawc2/shear_file.py b/wetb/hawc2/shear_file.py index 16c33f1a3243003e461b58e77eb24c3e3f3d85a4..882d83e78c7c33196268765a918cd0b2e6571cac 100644 --- a/wetb/hawc2/shear_file.py +++ b/wetb/hawc2/shear_file.py @@ -5,16 +5,46 @@ Created on 24/04/2014 ''' import numpy as np +import os +from wetb.functions.make_dirs import make_dirs def save(filename, y_coordinates, z_coordinates, u=None, v=None, w=None): + """ + Parameters + ---------- + filename : str + filename + y_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, #y_coordinates) or (#z_coordinates,) + v : array_like, optional + shear_v component, normalized with U_mean\n + shape must be (#z_coordinates, #y_coordinates) or (#z_coordinates,) + w : array_like, optional + shear_w component, normalized with U_mean\n + shape must be (#z_coordinates, #y_coordinates) or (#z_coordinates,) + """ + shape = (len(z_coordinates), len(y_coordinates)) vuw = [v, u, w] for i in range(3): if vuw[i] is None: - vuw[i] = np.zeros((shape)) + 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) + make_dirs(filename) with open(filename, 'w') as fid: fid.write(" # autogenerated shear file\n") fid.write(" %d %d\n" % (len(y_coordinates), len(z_coordinates))) diff --git a/wetb/hawc2/tests/test_shear_file.py b/wetb/hawc2/tests/test_shear_file.py index 4dbca846548a35d0ed6e5dd1933efd283b0411af..f5274bff2d0bb672c982df6b6f2e9ad584505b64 100644 --- a/wetb/hawc2/tests/test_shear_file.py +++ b/wetb/hawc2/tests/test_shear_file.py @@ -7,6 +7,7 @@ import unittest from wetb.hawc2 import shear_file import numpy as np import os +import shutil testfilepath = 'test_files/' class Test(unittest.TestCase): @@ -41,6 +42,40 @@ class Test(unittest.TestCase): os.remove(f) + def test_shearfile2(self): + f = testfilepath + "tmp_shearfile.dat" + shear_file.save(f, [-55, 55], [30, 100, 160] , u=np.array([0.7, 1, 1.3]).T) + with open(f) as fid: + self.assertEqual(fid.read(), +""" # autogenerated shear file + 2 3 + # shear v component + 0.0000000000 0.0000000000 + 0.0000000000 0.0000000000 + 0.0000000000 0.0000000000 + # shear u component + 0.7000000000 0.7000000000 + 1.0000000000 1.0000000000 + 1.3000000000 1.3000000000 + # shear w component + 0.0000000000 0.0000000000 + 0.0000000000 0.0000000000 + 0.0000000000 0.0000000000 + # y coordinates + -55.0000000000 + 55.0000000000 + # z coordinates + 30.0000000000 + 100.0000000000 + 160.0000000000 +""") + os.remove(f) + + def test_shear_makedirs(self): + f = testfilepath + "shear/tmp_shearfile2.dat" + shear_file.save(f, [-55, 55], [30, 100, 160] , u=np.array([0.7, 1, 1.3]).T) + shutil.rmtree(testfilepath + "shear") + if __name__ == "__main__": #import sys;sys.argv = ['', 'Test.test_shearfile'] unittest.main()