documentation + ability to take 1d u,v,w vectors

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)))

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()