def coriolis_parameter(lat,
gravity=default_gravity,
fromvar=False,
format_axes=False):
"""Get the coriolis parameters computed at each latitude
:Params:
- **lat**: Latitude or a variable with latitude coordinates.
- **gravity**, optional: Gravity.
- **fromvar**, optional: If True, lat is supposed to be a MV2
array with latitude coordinates.
"""
# Latitude
if fromvar:
if not cdms2.isVariable(lat):
raise VACUMMError('lat must a MV2 array because fromvar is True')
latv = lat * 0
lat = lat.getLatitude()
if lat is None:
raise VACUMMError(
'lat must a MV2 array with a latitude axis because fromvar is True'
)
if cdms2.isVariable(lat): lat = lat.asma() # 2D axes
if lat.shape != latv.shape:
if len(lat.shape) == 2:
latv[:] = N.ma.resize(lat, latv.shape)
else:
yaxis = latv.getOrder().index('y')
new_shape = len(latv.shape) * [1]
new_shape[yaxis] = latv.shape[yaxis]
tile_shape = list(latv.shape)
tile_shape[yaxis] = 1
latv[:] = N.tile(lat[:].reshape(new_shape), tile_shape)
else:
latv = lat if not N.ndim(lat) else lat[:]
# Compute
f0 = 2 * N.ma.sin(N.pi * latv / 180.)
# f0 *= 2*N.pi/(24.*3600.)
f0 *= 2 * N.pi / (86164.) # 86164 = sidereal day....
# Format
if N.isscalar(f0): return f0
f0 = MV2.asarray(f0)
if not fromvar and isaxis(lat) and f0.ndim == 1:
f0.setAxis(0, lat)
return format_var(f0, 'corio', format_axes=format_axes)
def get_data_dir(raiseerr=False):
"""Get the data directory absolute path
This directory contains data samples and other needed files.
It can be at two different places, depending on if the library is
an installed version or a developers version.
- If installed : in :file:`vacumm-data` subdirectory in the
installed package directory (see :meth:`get_lib_dir`).
- Else in the :file:`data` subdirectory of the main distribution tree
(see :meth:`get_dist_dir`).
.. warning:: It raises an :class:`VACUMMError` error if not found
and if ``raiseerr`` is True, else return ''.
"""
# Installed librairy
lib_dir = get_lib_dir()
data_dir = os.path.join(lib_dir, 'vacumm-data')
if os.path.exists(data_dir):
return data_dir
# Distributed library (dev)
dist_dir = get_dist_dir()
if dist_dir is not None:
data_dir = os.path.join(dist_dir, 'data')
if os.path.exists(data_dir):
return data_dir
if raiseerr: raise VACUMMError("Can't find a valid data directory")
return ''
def density(temp,
sal,
depth=None,
lat=None,
potential=False,
getdepth=False,
getlat=False,
format_axes=False):
"""Compute density from temperature, salinity and depth (and latitude)
:Params:
- **temp**: Insitu or potential temperature.
- **sal**: Salinity.
- **depth**, optional: Depth at temperature and salinty points.
Assumed to be 0 if not found.
- **lat**, optional: Latitude. Error when not found.
- **potential**, optional: True to get the potential density (at atmospheric
pressure).
:Algo:
>>> pressure = seawater.csiro.pres(depth, lat)
>>> density = seawater.csiro.dens(sal, temp, depth)
"""
# Compute
if not potential and depth is not False: # In-situ
# Get depth and latitude
lat = grow_lat(temp, lat, mode='raise', getvar=False)
if lat is None: raise VACUMMError('No latitude found for density')
depth = grow_depth(temp, depth, mode='raise', getvar=False)
if N.abs(depth.max()) < N.abs(depth.min()): # positive
depth = -depth
if (depth.asma() < 0).any():
depth = depth - depth.min() # top=0
# Get density
pres = sw_pres(depth, lat)
dens = sw_dens(sal, temp, pres)
del pres
else: # Potential
dens = sw_dens0(sal, temp)
getdepth = getlat = False
# Format
dens.setAxisList(temp.getAxisList())
set_grid(dens, get_grid(temp))
format_var(dens, 'dens', format_axes=format_axes)
# Out
if not getdepth and not getlat: return dens
dens = dens,
if getdepth: dens += depth,
if getlat: dens += lat,
return dens
def get_scripts_dir(subdir=None, raiseerr=False):
"""Get the scripts directory absolute path
This directory contains examples of script.
It can be at two different places, depending on if the library is
an installed version or a developers version.
- If installed : in :file:`vacumm-scripts` subdirectory in the
installed package directory (see :meth:`get_lib_dir`).
- Else in the :file:`scripts` subdirectory of the main distribution tree
(see :meth:`get_dist_dir`).
.. warning:: It raises an :class:`VACUMMError` error if not found
and if ``raiseerr`` is True.
"""
# Installed librairy
lib_dir = get_lib_dir()
scripts_dir = os.path.join(lib_dir, 'vacumm-scripts')
if not os.path.exists(scripts_dir): scripts_dir = None
# Distributed library (dev)
if scripts_dir is None:
dist_dir = get_dist_dir()
if dist_dir is not None:
scripts_dir = os.path.join(dist_dir, 'scripts')
if not os.path.exists(scripts_dir): scripts_dir = None
# Not found
if scripts_dir is None:
if raiseerr: raise VACUMMError("Can't find a valid scripts directory")
return ''
# Subdir
if subdir and isinstance(subdir, basestring):
scripts_dir = os.path.join(scripts_dir, subdir)
if not os.path.exists(scripts_dir):
raise VACUMMError(
"Invalid subdirectory of the scripts directory: " + subdir)
return scripts_dir
def cf2search(name, mode='isa', raiseerr=True, **kwargs):
"""Extract specs from :attr:`CF_AXIS_SPECS` or :attr:`CFVAR_SPECS`
to form a search dictionary
:Params:
- **name**: Generic name of an axis or a variable.
- **mode**, optional: Search mode [default: None->``"ns"``].
A string containg one or more of the following letters:
- ``"n"``: Search using names (ids).
- ``"s"``: Search using standard_name attribute.
- ``"a"``: Search using axis attribute.
- ``"l"``: Search using long_name attribute.
- ``"u"``: Search using units attribute.
The order is important.
:Return: An :class:`colections.OrderedDict`
:Example:
>>> cf2search('sst', mode='isu')
{'id':['sst'], 'id':['sst'],
'standard_names':['sea_surface_temperature'],
'units':['degrees_celsius']}
"""
# Get specs
if name in CF_VAR_SPECS:
specs = CF_VAR_SPECS[name]
elif name in CF_AXIS_SPECS:
specs = CF_AXIS_SPECS[name]
else:
if raiseerr:
raise VACUMMError("Wrong generic name. It should be one of: " +
' '.join(CF_AXIS_SPECS.keys() +
CF_VAR_SPECS.keys()))
else:
return
# Form search dict
if not isinstance(mode, basestring):
mode = 'isa'
mode = mode.replace('n', 'i')
keys = []
for m in mode:
for key in ['id', 'standard_name', 'axis', 'long_name', 'units']:
if key.startswith(m):
keys.append(key)
break
return OrderedDict([(k, specs[k]) for k in keys if k in specs])
def DS(ncfile, clsname='generic', *args, **kwargs):
"""Load a specialized :class:`vacumm.data.misc.dataset.Dataset` instance
Available methods for retreiving data (or derived data) are the same
for all dataset types.
And whatever the dataset type, the outputs will all have the same
format (names, attributes...). Formating is defined and performed
by the :mod:`vacumm.data.cf` module.
:Params:
- **ncfile**: Netcdf file name(s) or pattern compatibile with
:func:`~vacumm.misc.io.list_forecast_files`, so
please read the documentation of this function carefully.
- **clsname**: Generic class name. Please choose one of:
%s
- All other keyword are passed to the class initialization.
See :class:`~vacumm.data.misc.dataset.Dataset` a list of these options.
:Return: A child of :class:`vacumm.data.misc.dataset.Dataset`.
:Example:
>>> from vacumm.data import DS
>>> mars = DS('results.nc', 'mars')
>>> dx,dy = mars.get_resol()
>>> sst = mars.get_sst()
:See also: :ref:`user.desc.dataset`.
"""
# Register builtin datasets
import vacumm.data.misc.dataset
import vacumm.data.model
# Class name
if clsname is None:
clsname = 'generic'
clsname = clsname.lower()
# Class
if clsname in DATASET_SPECS:
cls = DATASET_SPECS[clsname]
else:
raise VACUMMError('Wrong name of dataset type: %s. '
' Please choose one of the following: %s' %
(clsname, ', '.join(DATASET_NAMES)))
# Instantiate
return cls(ncfile, *args, **kwargs)
def cf2atts(name, select=None, exclude=None, ordered=True, **extra):
"""Extract specs from :attr:`CF_AXIS_SPECS` or :attr:`CF_VAR_SPECS` to form
a dictionary of attributes (units and long_name)"""
# Get specs
if isinstance(name, dict):
specs = name.copy()
elif name in CF_VAR_SPECS:
specs = CF_VAR_SPECS[name]
elif name in CF_AXIS_SPECS:
specs = CF_AXIS_SPECS[name]
else:
raise VACUMMError("Wrong generic name: %s. It should be one of: " %
name +
' '.join(CF_AXIS_SPECS.keys() + CF_VAR_SPECS.keys()))
# Which attributes
atts = OrderedDict() if ordered else {}
if exclude is None: exclude = []
elif isinstance(exclude, basestring): exclude = [exclude]
exclude.extend(_attnames_exclude)
for key in _attnames_firsts + specs.keys():
# Skip aliases
if key in BaseSpecs.get_alias_list():
continue
# Skip some attributes
if (key not in specs or key in exclude or key in atts
or (select is not None and key not in select)):
continue
# No lists or tuples
value = specs[key]
if isinstance(value, (list, tuple)):
if len(value) == 0:
continue
value = value[0]
# Store it
atts[key] = value
# Extra
for att, val in extra.items():
atts[att] = val
return atts
def register_dataset(cls, clsname=None, warn=True, force=True):
"""Register a new :class:`~vacumm.data.misc.dataset.Dataset` class
that can be loaded with :func:`DS`
"""
# Get the class when a path is given
if isinstance(cls, str):
modname, Clsname = _os.path.splitext(cls)
Clsname = Clsname[1:]
mod = __import__(modname, fromlist=[Clsname])
cls = getattr(mod, Clsname)
# Check inheritance
from vacumm.data.misc.dataset import Dataset
if not issubclass(cls, Dataset):
raise VACUMMError('cls must be a Dataset subclass')
# Get the name
if clsname is None:
clsname = cls.name
if clsname is None:
clsname = cls.__name__.lower()
clsname = clsname.lower()
if cls.name is None:
cls.name = clsname
# Already registered?
if clsname in DATASET_SPECS:
if warn:
if force:
ms = 'Overwriting it...'
else:
ms = 'Skipping...'
sonat_warn('Dataset class "{}" is already registered. ' + ms)
if not force:
return
# Register
DATASET_SPECS[clsname] = cls
return cls
def edit_file(fname, editor=None, verbose=True):
"""Edit a file
:Params:
- **fname**: File name.
- **editor**, optional: Editor to use (default to :envvar:`VISUAL` or
:envvar:`EDITOR`.
:Source: Mercurial code
"""
# File name
if not os.path.exists(fname):
raise IOError('File not found: ' + fname)
# Guess editor
if editor is None:
editor = os.environ.get("VISUAL") or os.environ.get("EDITOR", "vi")
# Run editor
if verbose: print('Editing user configuration file: ' + fname)
if subprocess.call(shlex.split("%s \"%s\"" % (editor, fname))):
raise VACUMMError("Error editing file: " + fname)
if verbose: print('End of file edition')
def check_data_file(section,
option,
parent_section=None,
parent_option=None,
quiet=None,
suffix=None,
avail=False,
check_license=True):
"""Check the existence of a file whose path is stored in a configuration file
Two cases are treated:
1. Path value is accessible from the vacumm configuration (:func:`get_config`),
using ``section`` (module name) and ``option``.
2. Path value is accessible from a secondary configuration file using
``parent_section`` (module name) and ``parent_option``, and the
path value of this config file is accessible from the vacumm configuration
using ``section`` and ``name``.
If the file is not found, it may download it using an url whose value is
accessible at the same place as the path and with the same option
name + '_url'.
:Tasks:
#. If ``suffix`` is a list, call itself in a loop with ``suffix``
set to each element of the list.
#. Get the ``path``, ``url`` and ``license`` of the data file.
#. If ``avail is True``, only check the existence and return the status.
#. If the file is not found (and ``quiet is False``), ask the user
for it. If the path specified is empty, simply go further.
If is not empty, update the configuration and return
it, else raise an error.
#. If a license info is present, tell the user that he must
have the authorization to download it. If the user does not agree,
raise an error.
#. Ask the user where to download the file using :func:`get_dl_dir`.
#. Download the file using :func:`download_file`.
#. Update the configuration if needed.
#. Return the path.
:Params:
- **section**: Section where to find the path to data.
- **option**: Option to read the path.
- **parent_section**, optional: Section of the vacumm configuration
where to find the path to the secondary configuration file that has
the path stored within.
- **parent_option**, optional: Option of the vacumm configuration
to read the path to the secondary configuration.
- **quiet**, optional: Don't ask question about where to download the
file, don't display any info, don't ask for authorization.
- **suffix**, optional: A suffix or a list of suffixes to append to
the path and url before using it.
- **check_license**, optional: If a license info is found,
show ot and ask the user if he has autorization for downloading it.
- **avail**, optional: Check availability only (see below).
:Return:
- A single path or a list of paths.
- ``None`` or a list of them if nothing found.
- If ``avail=True``:
- ``0``: Data file not on disk and not downloadable
(however, user can specify the path manually).
- ``1``: File is not on disk, but downloadable.
- ``2``: File is on disk.
:Examples:
>>> check_data_file('vacumm.bathy.shorelines', 'shapefile_histolitt', suffix=['.shp', '.dbf'])
>>> check_data_file('etopo2', 'file', parent_section='vacumm.bathy.bathy',
... parent_option='cfgfile_gridded')
"""
if quiet is None:
quiet = not sys.stdin.isatty()
# Loop on suffixes
if isinstance(suffix, (list, tuple)):
paths = []
for i, suf in enumerate(suffix):
paths.append(
check_data_file(section,
option,
parent_section=parent_section,
parent_option=parent_option,
quiet=quiet,
suffix=suf,
check_license=i == 0))
return paths
elif not isinstance(suffix, basestring):
suffix = ''
# Get local and remote file names
path, cfg = get_config_value(section,
option,
umode='merge',
getcfg=True,
parent_section=parent_section,
parent_option=parent_option)
url = get_config_value(section, option + '_url', umode='merge', cfg=cfg)
license = get_config_value(section,
option + '_license',
umode='merge',
cfg=cfg)
if path is None:
raise VACUMMError(
"Can't determine path of data file. Here are the config specs:\n"
"section='%(section)s', option='%(option)s', parent_section='%(parent_section)s', parent_option='%(parent_option)s'"
% locals())
path += suffix
if url is not None: url += suffix
# Only check availability
if os.path.exists(path):
if avail: return 2
return path
if avail:
return 0 if url is None else 1
# Ask for this path
if not quiet:
print("Can't find data file: " + path)
guessed = input(
"If you know what is this file and where it is on your system,\n"
"please enter its name here, or leave it empty: \n").strip()
if guessed:
if not os.path.exists(guessed):
print('File not found')
else:
_set_config_path_(guessed, section, option, parent_section,
parent_option, parent and cfgpath, suffix)
return guessed
# Check url
if not url:
if quiet: return
raise VACUMMError(
'Data file not found and not url provided for downloading it: ' +
path)
# License for downloading
if license and check_license:
nc = len(license) + 4
lic = '#' * nc + '\n'
lic += '# %s #\n' % license
lic += '#' * nc
print("VACUMM is about to download this data file: %s\n"%url + \
"We suppose you have requested the authorization and are not responsible for your choice.\n" + \
"If you're not sure you are allowed to do it, please abort.\n" + \
"For more information about this data file and associated distribution license, check this:\n"+lic)
while True:
try:
c = input("Would you like to download it? [y/N]\n")
except:
c = 'n'
if not c: c = 'n'
if c.startswith('n'):
raise VACUMMError(
"Download interrupted -> can't access to data")
if c.startswith('y'): break
# Download directory
dl_dir = path and os.path.dirname(path)
if not path or not os.access(dl_dir, os.W_OK | os.R_OK | os.X_OK):
dl_dir = get_dl_dir(quiet=quiet, suggest=dl_dir)
# Download
basename = url.split('/')[-1]
dl_path = os.path.join(dl_dir, basename)
download_file(url, dl_path, quiet=quiet)
# Fix configuration
if path != dl_path:
_set_config_path_(dl_path, section, option, parent_section,
parent_option, parent_section and cfgpath, suffix)
return dl_path
def set_config_value(section, option, value=None, quiet=False, cfgfile=None):
"""Create or update user configuration file
The directory of this file is given by :func:`get_user_conf_dir`
and its name is :file:`vacumm.cfg`. Therefore it should be:
:file:`~/.config/vacumm/vacumm.cfg`.
:Params:
- **section**: A module or its name.
- **option**: Option name.
- **value**: Value of the option. If ``None``, option is removed.
- **cfgfile**, optional: Configuration file. If not provided,
internal user configuration file is used
(:file:`vacumm.cfg` in user config directory -- given by :func:`get_user_conf_dir`).
:Example:
>>> set_config_value('vacumm.bathy.bathy', 'cfgfile_gridded',
'%(mod_dir)s/bathy.gridded.cfg') # set
>>> import vacumm.bathy.bathy as B
>>> set_config_value(B, 'cfgfile_gridded') # remove
:Output: Path of the configuration file
"""
if hasattr(section, '__name__'):
section = section.__name__
# Load or check
if cfgfile is None:
cfgfile = os.path.join(get_user_conf_dir(), 'vacumm.cfg')
if not os.access(cfgfile, os.W_OK):
if os.path.exists(cfgfile):
raise VACUMMError("Can't write to config file: " + cfgfile)
else:
udir = os.path.dirname(cfgfile)
try:
os.makedirs(udir)
except:
raise VACUMMError("Can't write to config file: " + cfgfile)
cfg = SafeConfigParser()
cfg.read(cfgfile)
# Update
if value is not None:
value = str(value)
if not cfg.has_section(section):
cfg.add_section(section)
cfg.set(section, option, value)
if not quiet:
print('Updated user configuration file (%s) with:' % cfgfile)
print(' [%(section)s]\n %(option)s=%(value)s' % locals())
elif cfg.has_option(section, option):
cfg.remove_option(section, option)
if not quiet:
print(
'Removed the following option from user configuration file (%s):'
% cfgfile)
print(' [%(section)s]\n %(option)s' % locals())
# Save
f = open(cfgfile, 'w')
cfg.write(f)
f.close()
return cfgfile
def format_var(var,
name=None,
force=True,
format_axes=True,
order=None,
nodef=True,
mode='warn',
**kwargs):
"""Format a MV2 variable according to its generic name
:Params:
- **var**: A :mod:`numpy` or :mod:`MV2` variabe.
- **name**: Generic name of variable. It should be one of
those listed by :attr:`CF_VAR_SPECS`. If None, it is guessed
with :func:`match_known_var`.
- **force**, optional: Overwrite attributes in all cases.
- **format_axes**, optional: Also format axes.
- **nodef**, optional: Remove location specification when it refers to the
default location (:attr:`DEFAULT_LOCATION`).
- **mode**: "silent", "warn" or "raise".
- Other parameters are passed as attributes, except those:
- present in specifications to allow overriding defaults,
- starting with 'x', 'y', or 't' which are passed
to :func:`format_axis`.
:Examples:
>>> var = format_var(myarray, 'sst', valid_min=-2, valid_max=100)
"""
# Filter keywords for axis formating
axismeths = {'t': 'getTime', 'y': 'getLatitude', 'x': 'getLongitude'}
kwaxes = {}
for k in axismeths.keys():
kwaxes[k] = kwfilter(kwargs, k + '_')
# Always a MV2 array
if not cdms2.isVariable(var):
var = MV2.asarray(var)
# Check specs
if name is None: # guess it
name = match_known_var(var)
if not name:
if mode == 'warn':
warn("Can't guess cf name")
return var
elif mode == 'silent':
return var
else:
raise KeyError("Variable does not match any CF var")
elif name not in CF_VAR_SPECS and name not in CF_AXIS_SPECS:
if var.id in CF_VAR_SPECS or var.id in CF_AXIS_SPECS:
name = var.id
elif mode == 'warn':
warn("Generic var name not found '%s'." % name)
return var
elif mode == 'silent':
return var
else:
raise KeyError(
"Generic var name not found '%s'. Please choose one of: %s" %
(name, ', '.join(CF_VAR_SPECS.keys() + CF_AXIS_SPECS.keys())))
isaxis = name in CF_AXIS_SPECS
if isaxis:
specs = CF_AXIS_SPECS[name].copy()
if 'axis' in specs:
del specs['axis']
else:
specs = CF_VAR_SPECS[name].copy()
# - merge kwargs and specs
for key, val in kwargs.items():
if val is None or key not in specs: continue
# Check type
if not isinstance(val, list) and isinstance(specs[key], list):
val = [val]
# Set
specs[key] = val
del kwargs[key]
# - remove default location
if nodef:
refloc = specs.get('physloc', None) or DEFAULT_LOCATION
for att in 'id', 'long_name', 'standard_name':
if get_loc(specs[att], att) == refloc:
specs[att] = [no_loc_single(specs[att][0], att)]
name = specs['id'][0]
# - id
if ((force is True or force in [2, 'id', 'all'])
or var.id.startswith('variable_')
or (isaxis and var.id.startswith('axis_'))): # FIXME: use regexp
var.id = name
# - attributes
forceatts = (force is True or force in ['atts', 'all']
or (isinstance(force, int) and force > 0))
for att, val in cf2atts(specs, **kwargs).items():
if forceatts or not getattr(var, att, ''):
setattr(var, att, val)
# - physical location
loc = get_loc(var, mode='ext')
if not loc and 'physloc' in specs:
loc = specs['physloc']
if loc:
if 'physloc' in specs and loc == specs['physloc']:
var._vacumm_cf_physloc = loc.lower()
set_loc(var, loc)
# - store cf name
var._vacumm_cf_name = name
# Axes
if format_axes:
# Order
order = var.getOrder() if not isinstance(order, basestring) else order
if order is not None:
if not re.match('^[xyzt-]+$', order):
raise VACUMMError("Wrong cdms order type: " + order)
if len(order) != var.ndim:
raise VACUMMError(
"Cdms order should be of length %s instead of %s" %
(var.ndim, len(order)))
# First check
if 'axes' in specs:
axspecs = specs['axes']
formatted = []
for key, meth in axismeths.items():
axis = getattr(var, meth)()
if order is not None: order.replace(key, '-')
if axis is not None:
format_axis(axis, axspecs[key], **kwaxes[key])
formatted.append(key)
# Check remaining simple axes (DOES NOT WORK FOR 2D AXES)
if order is not None and order != '-' * len(order):
for key in axismeths.keys():
if key in order and key not in formatted:
axis = var.getAxis(order.index(key))
format_axis(axis, axspecs[key], **kwaxes[key])
return var
def mixed_layer_depth(data,
depth=None,
lat=None,
zaxis=None,
mode=None,
deltatemp=.2,
deltadens=.03,
kzmax=0.0005,
potential=True,
format_axes=False):
"""Get mixed layer depth from temperature and salinity
:Params:
- **temp**: Insitu or potential temperature.
- **sal**: Salinity.
- **depth**, optional: Depth at temperature and salinty points.
- **lat**, optional: Latitude.
- **mode**, optional: ``"deltatemp"``, ``"deltadens"``, ``"kz"``
or ``"twolayers"``
:Raise: :class:`~vacumm.VACUMMError` if can't get depth (and latitude for density).
"""
# TODO: positive up
# Inspection
if isinstance(data, tuple): # data = temp,sal
temp, sal = data
# Get density
if mode != 'deltatemp':
res = density(temp,
sal,
depth=depth,
lat=lat,
format_axes=False,
potential=potential,
getdepth=True)
if isinstance(res, tuple):
dens, depth = res
else:
dens = res
dens = dens.asma()
if mode is None:
mode = 'deltadens'
else:
temp = data[0]
# Check mode
if mode == 'kz':
warn("Switching MLD computation mode to 'deltadens'")
mode = "deltadens"
elif match_var(data, 'temp', mode='nslu'):
if mode is not None and mode != 'deltatemp':
warn("Switching MLD computation mode to 'deltatemp'")
mode = 'deltatemp'
temp = data
elif match_var(data, 'dens', mode='nslu'):
if mode in ['kz', 'deltatemp']:
warn("Switching MLD computation mode to 'deltadens'")
mode = None
if mode is None:
mode = "deltadens"
dens = data
elif match_var(data, 'kz', mode='nslu'):
if mode is None:
mode = "kz"
if mode != "kz":
warn("Switching MLD computation mode to 'kz'")
kz = data
else:
if mode in ['deltadens', 'twolayers']:
dens = data
elif mode == "deltatemp":
temp = data
elif mode == "kz":
kz = data
elif mode is not None:
raise VACUMMError("Invalid MLD computation mode : '%s'" % mode)
else:
raise VACUMMError("Can't guess MLD computation mode")
temp = delta
# Find Z dim
data0 = data[0] if isinstance(data, tuple) else data
depth = grow_depth(data0, depth, mode='raise', getvar=False)
zaxis = get_zdim(data0, axis=zaxis)
if zaxis is None:
raise VACUMMError("Can't guess zaxis")
slices = get_axis_slices(data0, zaxis)
# Init MLD
axes = data0.getAxisList()
del axes[zaxis]
mld = MV2.array(data0.asma()[slices['first']],
copy=1,
axes=axes,
copyaxes=False)
set_grid(mld, get_grid(data0))
format_var(mld, 'mld', format_axes=format_axes)
mld[:] = MV2.masked
# Two-layers
if mode == 'twolayers':
densbot = dens[slices['first']]
denstop = dens[slices['last']]
del dens
H = 1.5 * depth[slices['first']] - 0.5 * depth[slices['firstp1']]
H = -1.5 * depth[slices['last']] + 0.5 * depth[slices['lastm1']]
mld[:] = -H * (densbot - denstop) / (densbot - denstop)
del H
elif mode == 'deltadens':
denscrit = dens[slices['last']] + deltadens
mld[:] = -_val2z_(dens, depth, denscrit, zaxis, -1)
del dens
elif mode == 'deltatemp':
tempcrit = temp[slices['last']] - deltatemp
mld[:] = -_val2z_(temp, depth, tempcrit, zaxis, 1)
elif mode == 'kz':
mld[:] = -_valmin2z_(kz, depth, kzmax, zaxis, 1)
else:
raise VACUMMError("Invalid mode for computing MLD (%s)." % mode +
"Please choose one of: deltadens, twolayers")
# Mask zeros
mld[:] = MV2.masked_values(mld, 0., copy=0)
return mld
def check_order(var,
allowed,
vertical=None,
copy=False,
reorder=False,
extended=None,
getorder=False):
"""Check that the axis order of a variable is matches
at least one the specifed valid orders
:Params:
- **var**: MV2 array.
- **allowed**: A single order string or a list. It should contain one or
several of these letters:
- ``x``: longitudes,
- ``y``: latitudes,
- ``z``: vertical levels,
- ``t``: time axis,
- ``d``: data values (ignored),
- ``-``: any kind of axis.
:Return:
``var``, or ``var, order, reordered`` if **reorder** is True.
"""
# Check allowed orders
# - consistency
if not isinstance(allowed, (list, tuple)):
allowed = [allowed]
else:
allowed = list(allowed)
withd = 'd' in allowed[0]
get_rank = lambda o: len(o.replace('d', ''))
rank = get_rank(allowed[0])
for order in allowed:
try:
cdms2.orderparse(order.lower().replace('d', ''))
except:
raise VACUMMError("Wrong allowed order: " + order)
if ('d' in order and not withd) or ('d' not in order and withd):
raise VACUMMError(
"'d' only partially present in allowed order: %s" % allowed)
if get_rank(order) != rank:
raise VACUMMError("Inconsistent ranks between allowed orders: %s" %
[get_rank(o) for o in allowed])
# - check extended mode
if extended is None: # extended?
re_upper = re.compile('[XYZT]').search
for order in allowed:
if re_upper(order) is not None:
extended = True # force extended mode
break
else:
extended = False
if extended is False: # lower
allowed = [order.lower() for order in allowed]
else: #add tolerance for lower case orders
re_sub = re.compile('[xyzt]').sub
allowed = allowed + [re_sub('-', order) for order in allowed]
# - unique and lower case
_, idx = N.unique(allowed, return_index=True)
idx = N.sort(idx)
allowed = N.array(allowed)[idx].tolist()
# - restrict to vertical or horizontal (1D data only)
if vertical is not None and len(allowed[0]) == 2:
allowed = [oo for oo in allowed if oo[int(vertical)] == 'd']
# Data order
data_cdms_order = get_order(var)
# Loop on allowed orders
from vacumm.misc.grid.misc import get_axis, var2d
reordered = False
for allowed_order in allowed:
# Handle data case
d = allowed_order.find('d')
if d != -1: allowed_order = allowed_order.replace('d', '') # pure axis
# Check cdms order
allowed_cdms_order = allowed_order.lower() # lower case
if order_match(data_cdms_order, allowed_cdms_order, strict='right'):
break # It is already good
# Try to reorder
if reorder:
try:
reordered = cdms2.order2index(var.getAxisList(),
allowed_cdms_order)
new_var = var.reorder(allowed_cdms_order)
if allowed_cdms_order[-1] == 'x' and len(
get_axis(new_var, -1).shape) == 2: # 2D axes
del var
var = var2d(new_var,
MV2.transpose(get_axis(new_var, 0)),
MV2.transpose(get_axis(new_var, -1)),
copy=0)
set_order(new_var, allowed_cdms_order)
else:
del var
var = new_var
data_cdms_order = get_order(var)
break # No error so it worked and we leave
except:
continue
else:
raise VACUMMError('Wrong type of axes. Possible forms are: %s' %
', '.join(allowed))
if not getorder: return var
if d != -1:
data_cdms_order = cdms2.orderparse(data_cdms_order)
data_cdms_order.insert(d, 'd')
data_cdms_order = ''.join(data_cdms_order)
return var, data_cdms_order, reordered
def merge_orders(order1, order2, raiseerr=True):
"""Merge two axis orders
When two orders doesn't have the same length,
they are right adjusted.
:Examples:
>>> merge_orders('t-x', 'y-')
'tyx', 'yx'
>>> merge_orders('yx', 'tz')
'yx', 'tz'
>>> merge_orders(myvar, zaxis)
"""
order1 = get_order(order1)
order2 = get_order(order2)
rev = slice(None, None, 1 - 2 * int(len(order2) < len(order1)))
order1, order2 = (order1, order2)[rev]
# Inner loop
ishift = 0
n1 = len(order1)
n2 = len(order2)
n12 = n2 - n1
for i in range(n12 + 1):
j = n12 - i
if order_match(order1, order2[j:j + n1]):
i1 = 0
i2 = j
l = n1
break
else: # Outerloops
for ishift in range(1, min(n1, n2)):
l = min(n1, n2) - ishift
if order_match(order1[:l], order2[ishift:ishift + l]):
i1 = 0
i2 = ishift
break
if order_match(order2[:l], order1[ishift:ishift + l]):
i1 = ishift
i2 = 0
break
else:
if raiseerr:
raise VACUMMError(
'orders are incompatible and cannot be safely merged: %s %s'
% (order1, order2)[rev])
return (order1, order2)[rev]
# Merge
neworder1 = order1[:i1]
neworder2 = order2[:i2]
for i in range(l):
c1 = order1[i1 + i]
c2 = order2[i2 + i]
if c1 == c2 or c2 == '-':
neworder1 += c1
neworder2 += c1
elif c1 == '-':
neworder1 += c2
neworder2 += c2
else:
if raiseerr:
raise VACUMMError(
'orders are incompatible and cannot be safely merged: %s %s'
% (order1, order2)[rev])
return (order1, order2)[rev]
neworder1 += order1[i1 + l:]
neworder2 += order2[i2 + l:]
# Check multiples
for c in 'xyztd':
if neworder1.count(c) > 2 or neworder2.count(c) > 2:
warn('Merging of orders (%s and %s) may have not '%(order1, order2) + \
'properly worked (multiple axes are of the same type)')
return (neworder1, neworder2)[rev]