accelerate_workflow (#57)

* accelerate_workflow * squeeze axis and test * doc * requested changes
ocean-transport · Jun 1, 2017 · 87149bb · 87149bb
1 parent f771c0f
commit 87149bb
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 20 deletions.
diff --git a/floater/generators.py b/floater/generators.py
@@ -344,18 +344,17 @@ def npart_to_2D_array(self, ds1d):
         RETURNS
         -------
         ds2d : 2D Dataset
-            Two-dimensional dataset of physical variable(s) with dimensions 'lat' and 'lon'
+            Two-dimensional dataset of physical variable(s) with dimensions 'y0' and 'x0'
         """
 
         Nx = self.Nx
         Ny = self.Ny
         Nt = Nx*Ny
         if type(ds1d) == xr.core.dataarray.DataArray:
             ds1d = ds1d.to_dataset()
-        df = ds1d.to_dataframe()
-        var_list = list(df.columns)
-        index_dict = {'index': range(1, Nt+1)}
-        var_dict = {var: np.zeros(Nt) for var in var_list}
+        index_dict = {'index': np.arange(1, Nt+1, dtype=np.int32)}
+        var_list = list(ds1d.data_vars)
+        var_dict = {var: np.full(Nt, np.nan, dtype=np.float32) for var in var_list}
         frame_dict = {}
         frame_dict.update(index_dict)
         frame_dict.update(var_dict)
@@ -365,23 +364,27 @@ def npart_to_2D_array(self, ds1d):
         if self.model_grid is not None:
             ocean_bools = self.ocean_bools
         else:
-            ocean_bools = np.zeros(Nt, dtype=bool)==False
-        framei.loc[ocean_bools==True] = df.values.astype(np.float32)
-        framei.loc[ocean_bools==False] = np.float32('nan')
-        data_vars = {}
+            ocean_bools = np.full(Nt, True, dtype=bool)
+        da = ds1d.to_array().values
+        axis_len = len(da.shape) - 2
+        axis = tuple(np.arange(1, axis_len+1, dtype=np.int32))
+        das = np.squeeze(da, axis=axis)
+        dast = das.transpose()
+        framei.loc[ocean_bools==True] = dast.astype(np.float32)
         dim_list = list(ds1d.dims)
         dim_list.remove('npart')
         dim_len = len(dim_list)
         new_shape = (1,)*dim_len + (Ny, Nx)
-        new_dims = dim_list + ['lat', 'lon']
+        new_dims = dim_list + ['y0', 'x0']
+        data_vars = {}
         for var in var_list:
             frameir = framei[var].values
             frameir.shape = new_shape
             data_vars.update({var: (new_dims, frameir)})
+        x0 = np.float32(self.x)
+        y0 = np.float32(self.y)
         coords = {}
-        lon = np.float32(self.x)
-        lat = np.float32(self.y)
         coords.update({dim: ([dim], ds1d[dim].values) for dim in dim_list})
-        coords.update({'lat': (['lat'], lat), 'lon': (['lon'], lon)})
+        coords.update({'y0': (['y0'], y0), 'x0': (['x0'], x0)})
         ds2d = xr.Dataset(data_vars=data_vars, coords=coords)
         return ds2d
diff --git a/floater/test/test_generators.py b/floater/test/test_generators.py
@@ -237,13 +237,13 @@ def test_npart_to_2D_array():
         assert da2d.to_array().values.shape == (1, 1, 1, fs.Ny, fs.Nx)
         assert ds2d.to_array().values.shape == (3, 1, 1, fs.Ny, fs.Nx)
         # dimension test
-        assert da2d.dims == {'date': 1, 'loc': 1, 'lat': 9, 'lon': 9}
-        assert ds2d.dims == {'date': 1, 'loc': 1, 'lat': 9, 'lon': 9}
+        assert da2d.dims == {'date': 1, 'loc': 1, 'y0': 9, 'x0': 9}
+        assert ds2d.dims == {'date': 1, 'loc': 1, 'y0': 9, 'x0': 9}
         # coordinates test
-        np.testing.assert_allclose(da2d.lon.values, fs.x)
-        np.testing.assert_allclose(da2d.lat.values, fs.y)
-        np.testing.assert_allclose(ds2d.lon.values, fs.x)
-        np.testing.assert_allclose(ds2d.lat.values, fs.y)
+        np.testing.assert_allclose(da2d.x0.values, fs.x)
+        np.testing.assert_allclose(da2d.y0.values, fs.y)
+        np.testing.assert_allclose(ds2d.x0.values, fs.x)
+        np.testing.assert_allclose(ds2d.y0.values, fs.y)
         # values test
         da1d_values = values_list[0][0][0]
         da2d_values_full = da2d.to_array().values[0].ravel()

diff --git a/floater/test/test_utils.py b/floater/test/test_utils.py
@@ -115,7 +115,7 @@ def test_floats_to_netcdf(tmpdir, mitgcm_float_datadir_csv):
     mfdm = xr.open_mfdataset('test_netcdf/prefix_test.*.nc')
 
     # dimensions test
-    dims = [{'time': 2, 'npart': 40}, {'time': 2, 'lat': 4, 'lon': 10}]
+    dims = [{'time': 2, 'npart': 40}, {'time': 2, 'y0': 4, 'x0': 10}]
     assert mfdl.dims == dims[0]
     assert mfdm.dims == dims[1]