Skip to content

Commit

Permalink
Depth (#70)
Browse files Browse the repository at this point in the history 
* add depth dimension

* Updated to pass tests

* more test changes

* Improve if statements

* improve consistency of if statements

* Handle single-element zvect

* new tests for 3D

* Tweaks to function names
  • Loading branch information
@cspencerjones @rabernat
cspencerjones authored and rabernat committed Aug 31, 2018
1 parent 651436d commit 004fe5a
Show file tree
Hide file tree
Showing 2 changed files with 240 additions and 37 deletions.
130 changes: 93 additions & 37 deletions floater/generators.py
View file
Expand Up @@ -10,13 +10,14 @@


def geo_to_xyz(geo_cord):
lon_deg, lat_deg = geo_cord.transpose()
lon_deg, lat_deg,depth =np.vsplit(geo_cord,3)
earthrad=6371000
lat = np.radians(lat_deg)
lon = np.radians(lon_deg)
x = np.cos(lat)*np.cos(lon)
y = np.cos(lat)*np.sin(lon)
z = np.sin(lat)
xyz_cord = np.transpose(np.array((x,y,z)))
x = (earthrad+depth)*np.cos(lat)*np.cos(lon)
y = (earthrad+depth)*np.cos(lat)*np.sin(lon)
z = (earthrad+depth)*np.sin(lat)
xyz_cord = np.squeeze(np.transpose(np.array((x,y,z)),(2,0,1)))
return xyz_cord

def xyz_to_geo(xyz_coord):
Expand All @@ -32,7 +33,7 @@ def xyz_to_geo(xyz_coord):
class FloatSet(object):
"""Represents a set of initial float positions on a regular grid."""

def __init__(self, xlim=None, ylim=None, dx=1., dy=1., model_grid=None, load_path=None):
def __init__(self, xlim=None, ylim=None, dx=1., dy=1., model_grid=None, load_path=None, zvect=None):
"""Initialize FloatSet according to specified rectangular grid geometry.
PARAMETERS
Expand All @@ -41,6 +42,7 @@ def __init__(self, xlim=None, ylim=None, dx=1., dy=1., model_grid=None, load_pat
minimum and maximum x coordinate of cell vertices
ylim : tuple of two floats
minimum and maximum y coordinate of cell vertices
zvect : 1d array of depths
dx : float
grid spacing in x direction
dy : float
Expand All @@ -51,8 +53,11 @@ def __init__(self, xlim=None, ylim=None, dx=1., dy=1., model_grid=None, load_pat
2d array of dimensions in C order: shape==(len(lat), len(lon))
An element is True iff the corresponding tracer cell grid point
is unmasked (ocean)
OR
3d array of dimensions in C order shape==(len(depth),len(lat), len(lon))
'lon': 1d array of the mask grid longitudes
'lat': 1d array of the mask grid latitudes
'rc' : 1d array of mask depths
load_path : str
The filename to load a saved floatset object from
Expand All @@ -76,6 +81,14 @@ def __init__(self, xlim=None, ylim=None, dx=1., dy=1., model_grid=None, load_pat
self.Ny = int(self.Ly / self.dy)
self.x = self.xlim[0] + self.dx * np.arange(self.Nx) + self.dx/2
self.y = self.ylim[0] + self.dy * np.arange(self.Ny) + self.dy/2
if zvect is not None:
self.zvect = zvect
self.Nz=zvect.size
self.dims=3
else:
self.Nz=1
self.zvect=[-0.5]
self.dims=2
self.model_grid = model_grid
else:
self.from_pickle(load_path)
Expand All @@ -88,15 +101,23 @@ def get_rectmesh(self):
RETURNS
-------
x : np.ndarray
2D array of float x coordinates
1D array of float x coordinates
y : np.ndarray
2D array of float y coordinates
1D array of float y coordinates
z : np.ndarray
1D array of float z coordinates
"""

xx, yy = np.meshgrid(self.x, self.y)
if self.model_grid is not None:
xx, yy = self.subset_floats_from_mask(xx, yy)
return xx, yy
if self.dims==3:
xx, yy, zz= np.meshgrid(self.x, self.y,self.zvect)
if self.model_grid is not None:
xx, yy, zz = self.subset_floats_from_mask(xx, yy,zz)
return xx,yy,zz
else:
xx, yy = np.meshgrid(self.x, self.y)
if self.model_grid is not None:
zz=np.ones(xx.size)*-0.5
xx, yy = self.subset_floats_from_mask(xx, yy, zz)
return xx, yy


def get_hexmesh(self):
Expand All @@ -105,19 +126,30 @@ def get_hexmesh(self):
RETURNS
-------
x : np.ndarray
2D array of float x coordinates
1D array of float x coordinates
y : np.ndarray
2D array of float y coordinates
1D array of float y coordinates
z : np.ndarray
1D array of float z coordinates
"""
if self.dims==3:
xx, yy, zz = self.get_rectmesh()
# modify to be even-R horizontal offset
xx[::2] += self.dx/4
xx[1::2] -= self.dx/4
if self.model_grid is not None:
xx, yy, zz = self.subset_floats_from_mask(xx, yy, zz)
return xx,yy,zz
else:
xx, yy = self.get_rectmesh()
# modify to be even-R horizontal offset
xx[::2] += self.dx/4
xx[1::2] -= self.dx/4

xx, yy = self.get_rectmesh()
# modify to be even-R horizontal offset
xx[::2] += self.dx/4
xx[1::2] -= self.dx/4

if self.model_grid is not None:
xx, yy = self.subset_floats_from_mask(xx, yy)
return xx, yy
if self.model_grid is not None:
zz=np.ones(xx.size)*-0.5
xx, yy = self.subset_floats_from_mask(xx, yy, zz)
return xx, yy


def npart_index_to_ndarray(self, data, npart):
Expand All @@ -140,7 +172,10 @@ def parcel_area(self, latlon=True):
return self.dx * self.dy
else:
R = 6.371e6
lon, lat = self.get_rectmesh()
if self.dims==2:
lon, lat = self.get_rectmesh()
else:
lon, lat,depth = self.get_rectmesh()
dy = R * np.radians(self.dy)
dx = R * np.radians(self.dx) * np.cos(np.radians(lat))
# old code, wrong!
Expand Down Expand Up @@ -182,8 +217,12 @@ def to_mitgcm_format(self, filename, tstart=0, iup=0, mesh='rect',
raise ValueError('read_binary_prec should be 32 or 64; '
'got %g' % read_binary_prec )

if mesh == 'hex':
if mesh == 'hex' and self.dims==3:
xx, yy, zz = self.get_hexmesh()
elif mesh == 'hex' and self.dims==2:
xx, yy = self.get_hexmesh()
elif self.dims==3:
xx, yy, zz = self.get_rectmesh()
else:
xx, yy = self.get_rectmesh()
myx = xx
Expand All @@ -193,12 +232,16 @@ def to_mitgcm_format(self, filename, tstart=0, iup=0, mesh='rect',
# initial positions
lon = myx.ravel()
lat = yy.ravel()
if self.dims==2:
kpart=-0.5
else:
kpart=zz.ravel()

# other float properties

# kpart: depth of float release in meters, depth is negative, i.e. -1500
# for 1500 m
kpart = -0.5
#kpart = -0.5

# kfloat: target level of float (??)
kfloat = -0.5
Expand Down Expand Up @@ -267,7 +310,7 @@ def from_pickle(self, filename='./floatset.pkl'):



def subset_floats_from_mask(self, xx, yy):
def subset_floats_from_mask(self, xx, yy, zz):
"""Eliminate float positions that are on land land mask.
PARAMETERS
Expand All @@ -276,47 +319,59 @@ def subset_floats_from_mask(self, xx, yy):
float longitudes
yy : arraylike
float latitudes
zz : arraylike
float depths
model_grid : dictionary
the following key value pairs are expected
'land_mask': np.ndarray of bools
2d array of dimensions in C order: shape==(len(lat), len(lon))
An element is True iff the corresponding tracer cell grid point
is unmasked (ocean)
OR 3d array
'lon': 1d array of the mask grid longitudes
'lat': 1d array of the mask grid latitudes
'rc': 1d array of mask grid depths
RETURNS
-------
xx_masked, yy_masked: np.ndarray
xx_masked, yy_masked, zz_masked: np.ndarray
1D arrays of float coordinates subarrays:
"""

xx = xx.ravel()
yy = yy.ravel()
zz = zz.ravel()

mask_lon = self.model_grid['lon']
mask_lat = self.model_grid['lat']
land_mask = self.model_grid['land_mask']
if "rc" in self.model_grid:
mask_depth = self.model_grid['rc']
else:
mask_depth=[-0.5]
land_mask=np.expand_dims(land_mask,axis=2)
# we require the array to be using using C order
assert land_mask.shape == (len(mask_lat), len(mask_lon))
assert land_mask.shape[0:2] == (len(mask_lat), len(mask_lon))

# fast cartesian product
mask_geo = np.dstack(np.meshgrid(mask_lon, mask_lat))
mask_geo = mask_geo.reshape(-1, 2)
mask_geo = np.stack(np.meshgrid(mask_lon, mask_lat,mask_depth))
mask_geo = mask_geo.reshape(3,-1)

mask_bool_flat = land_mask.ravel()
mask_xyz = geo_to_xyz(mask_geo)
# a KDTree of the mask data in xyz form
mask_tree = cKDTree(mask_xyz)
floats_geo = np.transpose([xx.ravel(), yy.ravel()])
floats_geo = np.array([xx.ravel(), yy.ravel(), zz.ravel()])
# uniform hexagonal tiling
queries_xyz = geo_to_xyz(floats_geo)
# search for nearest neighbors
dist, neighbor_indices = mask_tree.query(queries_xyz, n_jobs=-1)
self.ocean_bools = np.take(mask_bool_flat, neighbor_indices.ravel()) # True -> neighbor is tracer ocean
floats_ocean = floats_geo[np.nonzero(self.ocean_bools.astype('int'))].T

return floats_ocean
floats_ocean = floats_geo[:,np.nonzero(self.ocean_bools.astype('int'))].T
floats_ocean=np.transpose(np.squeeze(floats_ocean))
if self.dims==2:
return floats_ocean[0:2,:]
else:
return floats_ocean


def npart_to_2D_array(self, ds1d):
Expand All @@ -335,7 +390,8 @@ def npart_to_2D_array(self, ds1d):

Nx = self.Nx
Ny = self.Ny
Nt = Nx*Ny
Nz=self.Nz
Nt = Nx*Ny*Nz
if type(ds1d) == xr.core.dataarray.DataArray:
ds1d = ds1d.to_dataset()
index_dict = {'index': np.arange(1, Nt+1, dtype=np.int32)}
Expand Down

0 comments on commit 004fe5a

Please sign in to comment.