# This file is part of OpenDrift.
#
# OpenDrift is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, version 2
#
# OpenDrift is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with OpenDrift. If not, see <https://www.gnu.org/licenses/>.
#
# Copyright 2015, Knut-Frode Dagestad, MET Norway
from collections import OrderedDict
import numpy as np
[docs]
class LagrangianArray:
"""A generic array-like class for Lagrangian particle tracking.
A LagrangianArray is a generic class keeping the values of given
properties ('variables') of a collection of particles at a given time.
Values are stored as named attributes (similar to recarray) which are
ndarrays (1D, vectors) with one value for each particle, or as scalars
for values shared among all particles.
This is an Abstract Base Class, meaning that only subclasses can be used.
Subclasses will add specific variables for specific purposes (particle
types, e.g. oil, fish eggs...) to the core variables described below.
Attributes:
variables: An OrderedDict where keys are names of the
variables/properties of the current object. The values
of the OrderedDict are dictionaries with names such as
'dtype', 'unit', 'standard_name' (CF), 'default' etc.
All variable names will be added dynamically as attributes of
the object after initialisation. These attributes will be
numpy ndarrays of same length, or scalars. The core variables
are:
- ID: an integer identifying each particle.
- status: 0 for active particles and a positive integer when deactivated
- lon: longitude (np.float32)
- lat: latitude (np.float32)
- z: vertical position of the particle in m, positive upwards (above sea surface)
"""
variables = OrderedDict([
('ID', {'dtype': np.int32, # Unique numerical identifier
'seed': False,
'default': -1}), # ID to be assigned by application
('status', {'dtype': np.int32, # Status categories
'seed': False,
'default': 0}),
('moving', {'dtype': np.int32, # Set to 0 for elements which are frosen
'seed': False,
'default': 1}),
('age_seconds', {'dtype': np.float32,
'units': 's',
'seed': False,
'default': 0}),
('origin_marker', {'dtype': np.int32,
'unit': '',
'description': 'An integer kept constant during the simulation. Different values may be used for different seedings, to separate elements during analysis. With GUI, only a single seeding is possible.',
'default': 0}),
('lon', {'dtype': np.float32,
'units': 'degrees_east',
'standard_name': 'longitude',
'long_name': 'longitude',
'seed': False,
'axis': 'X'}),
('lat', {'dtype': np.float32,
'units': 'degrees_north',
'standard_name': 'latitude',
'long_name': 'latitude',
'seed': False,
'axis': 'Y'}),
('z', {'dtype': np.float32,
'units': 'm',
'standard_name': 'z',
'long_name': 'vertical position',
'axis': 'Z',
'positive': 'up',
'default': 0})])
def __init__(self, **kwargs):
"""Initialises a LagrangianArray with given properties.
Args:
Keyword arguments (kwargs) with names corresponding to the
OrderedDict 'variables' of the class, and corresponding values.
The values must be ndarrays of equal length, or scalars.
All (or none) variables must be given, unless a default value
is specified in the OrderedDict 'variables'
An empty object may be created by giving no input.
"""
# Collect default values in separate dict, for easier access
default_values = {variable: self.variables[variable]['dtype'](
self.variables[variable]['default'])
for variable in self.variables
if 'default' in self.variables[variable]}
if len(kwargs) == 0:
# Initialise an empty object (all variables have length 0)
for var in self.variables:
kwargs[var] = []
# Check for missing arguments
missing_args = set(self.variables.keys()) - \
set(kwargs.keys()) - set(default_values.keys())
if missing_args:
raise TypeError('Missing arguments: %s' % str(list(missing_args)))
# Check for redundant arguments
redundant_args = set(list(kwargs.keys()) +
list(default_values.keys())) - set((self.variables.keys()))
if redundant_args:
raise TypeError('Redundant arguments: %s' %
str(list(redundant_args)))
# Check that input arrays have same length
array_lengths = [1]*len(kwargs)
for i, input_value in enumerate(kwargs.values()):
try:
array_lengths[i] = len(input_value)
except:
array_lengths[i] = 1 # scalar is given
if len(set(array_lengths) - {1}) > 1:
raise TypeError(
'Input arrays must have same length. Lengths given: '
+ str(array_lengths))
# Store input arrays
for default_variable in default_values.keys(): # set default values
setattr(self, default_variable, default_values[default_variable])
for input_variable in kwargs.keys(): # override with input values
setattr(self, input_variable, self.variables[input_variable]
['dtype'](kwargs[input_variable]))
# Store dtypes for all parameters in a common dtype object (for io)
self.dtype = np.dtype([(var[0], var[1]['dtype'])
for var in self.variables.items()])
# Status must always be array
if not type(self.status) == np.ndarray:
self.status = self.status*np.ones(self.lon.shape)
[docs]
@classmethod
def add_variables(cls, new_variables):
"""Method used by subclasses to add specific properties/variables."""
variables = cls.variables.copy()
variables.update(new_variables)
return variables
[docs]
def extend(self, other):
"""Add elements from another object."""
len_self = len(self)
len_other = len(other)
for var in self.variables:
present_data = getattr(self, var)
new_data = getattr(other, var)
# If both arrays have an identical scalar, it remains a scalar
if (not isinstance(new_data, np.ndarray) and
not isinstance(present_data, np.ndarray) and
present_data == new_data):
continue
else: # Otherwise we create arrays and concatenate
if not hasattr(present_data, '__len__'):
present_data = present_data*np.ones(len_self)
if not hasattr(new_data, '__len__'):
new_data = new_data*np.ones(len_other)
setattr(self, var, np.concatenate((present_data,
new_data)))
[docs]
def move_elements(self, other, indices):
"""Remove elements with given indices, and append to another object.
NB: indices is boolean array, not real indices!"""
# Move elements with given indices (boolean array)
# to another LagrangianArray
# NB: scalars and 1D arrays are converted to ndarrays and concatenated
self_len = len(self)
other_len = len(other)
for var in self.variables:
self_var = getattr(self, var)
other_var = getattr(other, var)
if (not isinstance(self_var, np.ndarray) and
not isinstance(other_var, np.ndarray)) and \
(other_var == self_var):
if np.sum(indices) == len(self):
setattr(self, var, []) # Empty if all elements moved
continue # Equal scalars - we do nothing
# Copy elements to other
self_var = np.atleast_1d(self_var)
other_var = np.atleast_1d(other_var)
if len(self_var) < self_len: # Convert scalar to array
self_var = self_var*np.ones(self_len)
if len(other_var) < other_len: # Convert scalar to aray
other_var = other_var*np.ones(other_len)
if len(self_var) > 0:
setattr(other, var, np.concatenate((other_var,
self_var[indices])))
else:
setattr(other, var, self_var[indices])
setattr(self, var, self_var[~indices]) # Remove from self
#if isinstance(self_var, np.ndarray) or\
# isinstance(other_var, np.ndarray): # Array
# setattr(other, var,
# np.concatenate((getattr(other, var),
# getattr(self, var)[indices])))
# # Remove elements from self
# setattr(self, var, getattr(self, var)[~indices])
#elif isinstance(getattr(other, var), np.ndarray):
# setattr(other, var,
# np.concatenate((getattr(other, var),
# np.atleast_1d(getattr(self, var)))))
#else:
# setattr(other, var, getattr(self, var)) # Scalar
[docs]
def __len__(self):
length = 0
for var in self.variables:
length = np.maximum(length, len(np.atleast_1d(getattr(self, var))))
return length
[docs]
def __repr__(self):
outStr = ''
for variable in self.variables.keys():
outStr += variable + ': ' + str(getattr(self, variable)) + '\n'
return outStr