ladim 1.3.3__py3-none-any.whl → 2.0.0__py3-none-any.whl

ladim/output.py

@@ -0,0 +1,246 @@
+from .model import Model, Module
+import netCDF4 as nc
+import numpy as np
+
+
+class Output(Module):
+    def __init__(self, model: Model):
+        super().__init__(model)
+
+
+class RaggedOutput(Output):
+    def __init__(self, model: Model, variables: dict, file: str, frequency):
+        """
+        Writes simulation output to netCDF file in ragged array format
+
+        :param model: Parent model
+        :param variables: Simulation variables to include in output, and their formatting
+        :param file: Name of output file, or empty if a diskless dataset is desired
+        :param frequency: Output frequency in seconds. Alternatively, as a two-element
+        tuple (freq_value, freq_unit) where freq_unit can be any numpy-compatible time
+        unit.
+
+        """
+        super().__init__(model)
+
+        # Convert output format specification from ladim.yaml config to OutputFormat
+        self._formats = {
+            k: OutputFormat.from_ladim_conf(v)
+            for k, v in variables.items()
+        }
+
+        self._init_vars = {k for k, v in self._formats.items() if v.is_initial()}
+        self._inst_vars = {k for k, v in self._formats.items() if v.is_instance()}
+
+        if not file:
+            from uuid import uuid4
+            self._fname = uuid4()
+            self._diskless = True
+        else:
+            self._fname = file
+            self._diskless = False
+
+        try:
+            freq_num, freq_unit = frequency
+        except TypeError:
+            freq_num = frequency
+            freq_unit = 's'
+        self._write_frequency = np.timedelta64(freq_num, freq_unit).astype('timedelta64[s]').astype('int64')
+
+        self._dset = None
+        self._num_writes = 0
+        self._last_write_time = np.int64(-4611686018427387904)
+
+    @property
+    def dataset(self) -> nc.Dataset:
+        """Returns a handle to the netCDF dataset currently being written to"""
+        return self._dset
+
+    def update(self):
+        if self._dset is None:
+            self._create_dset()
+
+        self._write_init_vars()
+        self._write_instance_vars()
+
+    def _write_init_vars(self):
+        """
+        Write the initial state of new particles
+        """
+
+        # Check if there are any new particles
+        part_size = self._dset.dimensions['particle'].size
+        num_new = self.model.state.released - part_size
+        if num_new == 0:
+            return
+
+        # Write variable data
+        idx = self.model.state['pid'] > part_size - 1
+        pid = self.model.state['pid'][idx]
+        for v in set(self._init_vars) - {'release_time'}:
+            # The idx array is not necessarily monotonically increasing by 1
+            # all the way. We therefore copy the data into a temporary,
+            # continuous array.
+            data_raw = self.model.state[v][idx]
+            data = np.zeros(num_new, dtype=data_raw.dtype)
+            data[pid - part_size] = data_raw
+            self._dset.variables[v][part_size:part_size + num_new] = data
+
+        # Write release time variable
+        data = np.broadcast_to(self.model.solver.time, shape=(num_new, ))
+        self._dset.variables['release_time'][part_size:part_size + num_new] = data
+
+    def _write_instance_vars(self):
+        """
+        Write the current state of dynamic varaibles
+        """
+
+        # Check if this is a write time step
+        current_time = self.model.solver.time
+        elapsed_since_last_write = current_time - self._last_write_time
+        if elapsed_since_last_write < self._write_frequency:
+            return
+        self._last_write_time = current_time
+
+        # Write current time
+        time_size = self._dset.dimensions['time'].size
+        time_value = current_time.astype('datetime64[s]').astype('int64')
+        self._dset.variables['time'][time_size] = time_value
+
+        # Write variable values
+        inst_size = self._dset.dimensions['particle_instance'].size
+        inst_num = self.model.state.size
+        inst_vars = {k: self.model.state[k] for k in set(self._inst_vars) - {'lat', 'lon'}}
+        if {'lat', 'lon'}.intersection(self._inst_vars):
+            x, y = self.model.state['X'], self.model.state['Y']
+            inst_vars['lon'], inst_vars['lat'] = self.model.grid.xy2ll(x, y)
+        for name, data in inst_vars.items():
+            self._dset.variables[name][inst_size:inst_size + inst_num] = data
+
+        # Write particle count
+        self._dset.variables['particle_count'][time_size] = inst_num
+
+    def _create_dset(self):
+        default_formats = dict(
+            time=OutputFormat(
+                ncformat='i8',
+                dimensions='time',
+                attributes=dict(
+                    long_name="time",
+                    standard_name="time",
+                    units="seconds since 1970-01-01",
+                ),
+            ),
+            instance_offset=OutputFormat(
+                ncformat='i8',
+                dimensions=(),
+                attributes=dict(long_name='particle instance offset for file'),
+            ),
+            particle_count=OutputFormat(
+                ncformat='i4',
+                dimensions='time',
+                attributes=dict(
+                    long_name='number of particles in a given timestep',
+                    ragged_row_count='particle count at nth timestep',
+                ),
+            ),
+            release_time=OutputFormat(
+                ncformat='i8',
+                dimensions='particle',
+                attributes=dict(
+                    long_name='particle release time',
+                    units='seconds since 1970-01-01',
+                )
+            )
+        )
+
+        self._dset = create_netcdf_file(
+            fname=self._fname,
+            formats={**default_formats, **self._formats},
+            diskless=self._diskless,
+        )
+
+        self._dset.variables['instance_offset'][:] = 0
+
+    def close(self):
+        if self._dset is not None:
+            self._dset.close()
+            self._dset = None
+
+
+class OutputFormat:
+    def __init__(self, ncformat, dimensions, attributes, kind=None):
+        self.ncformat = ncformat
+        self.dimensions = dimensions
+        self.attributes = attributes
+        self.kind = kind
+
+    def is_initial(self):
+        return self.kind == 'initial'
+
+    def is_instance(self):
+        return self.kind == 'instance'
+
+    @staticmethod
+    def from_ladim_conf(conf) -> "OutputFormat":
+        def get_keywords(ncformat='f4', kind='instance', **kwargs):
+            return dict(
+                props=dict(ncformat=ncformat, kind=kind),
+                attrs=kwargs,
+            )
+
+        keywords = get_keywords(**conf)
+        vkind = keywords['props']['kind']
+        if vkind == 'initial':
+            dims = 'particle'
+        elif vkind == 'instance':
+            dims = 'particle_instance'
+        else:
+            raise ValueError(f"Unknown kind: {vkind}")
+
+        return OutputFormat(
+            ncformat=keywords['props']['ncformat'],
+            dimensions=dims,
+            attributes=keywords['attrs'],
+            kind=vkind,
+        )
+
+
+def create_netcdf_file(fname: str, formats: dict[str, OutputFormat], diskless=False) -> nc.Dataset:
+    """
+    Create new netCDF file
+
+    :param fname: File name
+    :param formats: Formats, one entry for each variable
+    :param diskless: True if a memory dataset should be generated
+    :return: Empty, initialized dataset
+    """
+    from . import __version__ as ladim_version
+
+    dset = nc.Dataset(filename=fname, mode='w', format='NETCDF4', diskless=diskless)
+    dset.set_auto_mask(False)
+
+    # Create attributes
+    dset.Conventions = "CF-1.8"
+    dset.institution = "Institute of Marine Research"
+    dset.source = "Lagrangian Advection and Diffusion Model"
+    dset.history = "Created by ladim " + ladim_version
+    dset.date = str(np.datetime64('now', 'D'))
+
+    # Create dimensions
+    dset.createDimension(dimname="particle", size=None)
+    dset.createDimension(dimname="particle_instance", size=None)
+    dset.createDimension(dimname="time", size=None)
+
+    # Create variables
+    for varname, item in formats.items():
+        dset.createVariable(
+            varname=varname,
+            datatype=item.ncformat,
+            dimensions=item.dimensions,
+        )
+        dset.variables[varname].set_auto_mask(False)
+        dset.variables[varname].setncatts(item.attributes)
+
+
+    return dset

ladim/release.py

@@ -0,0 +1,234 @@
+import contextlib
+
+from .model import Model, Module
+import numpy as np
+import pandas as pd
+from .utilities import read_timedelta
+import logging
+
+
+logger = logging.getLogger(__name__)
+
+
+class Releaser(Module):
+    def __init__(self, model: Model):
+        super().__init__(model)
+
+
+class TextFileReleaser(Releaser):
+    def __init__(
+            self, model: Model, file, colnames: list = None, formats: dict = None,
+            frequency=(0, 's'), defaults=None,
+    ):
+        """
+        Release module which reads from a text file
+
+        The text file must be a whitespace-separated csv file
+
+        :param model: Parent model
+        :param file: Release file
+
+        :param colnames: Column names, if the release file does not contain any
+
+        :param formats: Data column formats, one dict entry per column. If any column
+        is missing, the default format is used. Keys should correspond to column names.
+        Values should be either ``"float"``, ``"int"`` or ``"time"``. Default value
+        is ``"float"`` for all columns except ``release_time``, which has default
+        value ``"time"``.
+
+        :param frequency: A two-element list with entries ``[value, unit]``, where
+        ``unit`` can be any numpy-compatible timedelta unit (such as "s", "m", "h", "D").
+
+        :param defaults: A dict of variables to be added to each particle. The keys
+            are the variable names, the values are the initial values at particle
+            release.
+        """
+
+        super().__init__(model)
+
+        # Release file
+        self._csv_fname = file   # Path name
+        self._csv_column_names = colnames   # Column headers
+        self._csv_column_formats = formats or dict()
+        self._dataframe = None
+
+        # Continuous release variables
+        self._frequency = read_timedelta(frequency)
+        self._last_release_dataframe = pd.DataFrame()
+        self._last_release_time = np.int64(-4611686018427387904)
+
+        # Other parameters
+        self._defaults = defaults or dict()
+
+    def update(self):
+        self._add_new()
+        self._kill_old()
+
+    def _kill_old(self):
+        state = self.model.state
+        if 'alive' in state:
+            alive = state['alive']
+            alive &= self.model.grid.ingrid(state['X'], state['Y'])
+            state.remove(~alive)
+
+    def _add_new(self):
+        # Get the portion of the release dataset that corresponds to
+        # current simulation time
+        df = release_data_subset(
+            dataframe=self.dataframe,
+            start_time=self.model.solver.time,
+            stop_time=self.model.solver.time + self.model.solver.step,
+        ).copy(deep=True)
+
+        # If there are no new particles and we don't use continuous release,
+        # we are done.
+        continuous_release = bool(self._frequency)
+        if (len(df) == 0) and not continuous_release:
+            return
+
+        # If we have continuous release, but there are no new particles and
+        # the last release is recent, we are also done
+        current_time = self.model.solver.time
+        elapsed_since_last_write = current_time - self._last_release_time
+        last_release_is_recent = (elapsed_since_last_write < self._frequency)
+        if continuous_release and (len(df) == 0) and last_release_is_recent:
+            return
+
+        # If we are at the final time step, we should not release any more particles
+        if continuous_release and self.model.solver.time >= self.model.solver.stop:
+            return
+
+        # If we have continuous release, but there are no new particles and
+        # the last release is NOT recent, we should replace the empty
+        # dataframe with the previously released dataframe
+        if continuous_release:
+            if (len(df) == 0) and not last_release_is_recent:
+                df = self._last_release_dataframe
+            self._last_release_dataframe = df  # Update release dataframe
+            self._last_release_time = current_time
+
+        # If positions are given as lat/lon coordinates, we should convert
+        if "X" not in df.columns or "Y" not in df.columns:
+            if "lon" not in df.columns or "lat" not in df.columns:
+                logger.critical("Particle release must have position")
+                raise ValueError()
+            # else
+            X, Y = self.model.grid.ll2xy(df["lon"].values, df["lat"].values)
+            df.rename(columns=dict(lon="X", lat="Y"), inplace=True)
+            df["X"] = X
+            df["Y"] = Y
+
+        # Add default variables, if any
+        for k, v in self._defaults.items():
+            if k not in df:
+                df[k] = v
+
+        # Expand multiplicity variable, if any
+        if 'mult' in df:
+            df = df.loc[np.repeat(df.index, df['mult'].values.astype('i4'))]
+            df = df.reset_index(drop=True).drop(columns='mult')
+
+        # Add new particles
+        new_particles = df.to_dict(orient='list')
+        state = self.model.state
+        state.append(new_particles)
+
+    @property
+    def dataframe(self):
+        @contextlib.contextmanager
+        def open_or_relay(file_or_buf, *args, **kwargs):
+            if hasattr(file_or_buf, 'read'):
+                yield file_or_buf
+            else:
+                with open(file_or_buf, *args, **kwargs) as f:
+                    yield f
+
+        if self._dataframe is None:
+            if isinstance(self._csv_fname, pd.DataFrame):
+                self._dataframe = self._csv_fname
+
+            else:
+                with open_or_relay(self._csv_fname, 'r', encoding='utf-8') as fp:
+                    self._dataframe = load_release_file(
+                        stream=fp,
+                        names=self._csv_column_names,
+                        formats=self._csv_column_formats,
+                    )
+        return self._dataframe
+
+
+def release_data_subset(dataframe, start_time, stop_time):
+    start_idx, stop_idx = sorted_interval(
+        dataframe['release_time'].values,
+        start_time,
+        stop_time,
+    )
+    return dataframe.iloc[start_idx:stop_idx]
+
+
+def load_release_file(stream, names: list, formats: dict) -> pd.DataFrame:
+    if names is None:
+        import re
+        first_line = stream.readline()
+        names = re.split(pattern=r'\s+', string=first_line.strip())
+
+    converters = get_converters(varnames=names, conf=formats)
+
+    df = pd.read_csv(
+        stream,
+        names=names,
+        converters=converters,
+        sep='\\s+',
+    )
+    df = df.sort_values(by='release_time')
+    return df
+
+
+def sorted_interval(v, a, b):
+    """
+    Searches for an interval in a sorted array
+
+    Returns the start (inclusive) and stop (exclusive) indices of
+    elements in *v* that are greater than or equal to *a* and
+    less than *b*. In other words, returns *start* and *stop* such
+    that v[start:stop] == v[(v >= a) & (v < b)]
+
+    :param v: Sorted input array
+    :param a: Lower bound of array values (inclusive)
+    :param b: Upper bound of array values (exclusive)
+    :returns: A tuple (start, stop) defining the output interval
+    """
+    start = np.searchsorted(v, a, side='left')
+    stop = np.searchsorted(v, b, side='left')
+    return start, stop
+
+
+def get_converters(varnames: list, conf: dict) -> dict:
+    """
+    Given a list of varnames and config keywords, return a dict of converters
+
+    Returns a dict where the keys are ``varnames`` and the values are
+    callables.
+
+    :param varnames: For instance, ['release_time', 'X', 'Y']
+    :param conf: For instance, {'release_time': 'time', 'X': 'float'}
+    :return: A mapping of varnames to converters
+    """
+    dtype_funcs = dict(
+        time=lambda item: np.datetime64(item, 's').astype('int64'),
+        int=int,
+        float=float,
+    )
+
+    dtype_defaults = dict(
+        release_time='time',
+    )
+
+    converters = {}
+    for varname in varnames:
+        dtype_default = dtype_defaults.get(varname, 'float')
+        dtype_str = conf.get(varname, dtype_default)
+        dtype_func = dtype_funcs[dtype_str]
+        converters[varname] = dtype_func
+
+    return converters

ladim/sample.py

@@ -277,6 +277,9 @@ def bilin_inv(f, g, F, G, maxiter=7, tol=1.0e-7):
             i = x.astype("i")
             j = y.astype("i")
 
+        i = np.maximum(np.minimum(i, imax - 2), 0)
+        j = np.maximum(np.minimum(j, jmax - 2), 0)
+
         p, q = x - i, y - j
 
         # Bilinear estimate of F[x,y] and G[x,y]

ladim/{timestepper.py → solver.py}

@@ -1,13 +1,13 @@
 import numpy as np
 
 
-class TimeStepper:
+class Solver:
     def __init__(self, modules, start, stop, step, order=None, seed=None):
-        self.order = order or ('forcing', 'release', 'output', 'ibm', 'tracker')
+        self.order = order or ('release', 'forcing', 'tracker', 'ibm', 'output')
         self.modules = modules
-        self.start = np.datetime64(start)
-        self.stop = np.datetime64(stop)
-        self.step = np.timedelta64(step, 's')
+        self.start = np.datetime64(start, 's').astype('int64')
+        self.stop = np.datetime64(stop, 's').astype('int64')
+        self.step = np.timedelta64(step, 's').astype('int64')
         self.time = None
 
         if seed is not None:

ladim/state.py

@@ -0,0 +1,142 @@
+import pandas as pd
+import numpy as np
+from .model import Model, Module
+
+
+class State(Module):
+    def __init__(self, model: Model):
+        """
+        The state module contains static and dynamic particle properties
+
+        The other modules interact with the state module mostly through
+        the getitem and setitem methods. For instance, to increase the
+        depth of all particles by 1, use
+
+        >>> model.state['Z'] += 1
+
+        :param model: Parent model
+        """
+        super().__init__(model)
+
+    @property
+    def size(self):
+        """
+        Current number of particles
+        """
+        raise NotImplementedError
+
+    @property
+    def released(self):
+        """
+        Total number of released particles
+        """
+        raise NotImplementedError
+
+    def append(self, particles: dict):
+        """
+        Add new particles
+
+        Missing variables are assigned a default value of 0.
+
+        :param particles: A mapping from variable names to values
+        """
+        raise NotImplementedError
+
+    def remove(self, particles):
+        """
+        Remove particles
+
+        :param particles: Boolean index of particles to remove
+        :return:
+        """
+        raise NotImplementedError
+
+    def __getitem__(self, item):
+        raise NotImplementedError
+
+    def __setitem__(self, key, value):
+        raise NotImplementedError
+
+    def __len__(self):
+        return self.size
+
+    def __contains__(self, item):
+        raise NotImplementedError
+
+
+class DynamicState(State):
+    def __init__(self, model: Model):
+        super().__init__(model)
+
+        self._num_released = 0
+        self._varnames = set()
+
+        self._data = pd.DataFrame()
+
+    @property
+    def released(self):
+        return self._num_released
+
+    def append(self, particles: dict):
+        # If there are no new particles, do nothing
+        if not particles:
+            return
+
+        num_new_particles = next(len(v) for v in particles.values())
+        particles['pid'] = np.arange(num_new_particles) + self._num_released
+        particles['alive'] = np.ones(num_new_particles, dtype=bool)
+        if 'active' in particles:
+            particles['active'] = np.array(particles['active'], dtype=bool)
+        else:
+            particles['active'] = np.ones(num_new_particles, dtype=bool)
+
+        new_particles = pd.DataFrame(data=particles)
+        self._data = pd.concat(
+            objs=[self._data, new_particles],
+            axis='index',
+            ignore_index=True,
+            join='outer',
+        )
+
+        self._num_released += num_new_particles
+
+    def remove(self, particles):
+        if not np.any(particles):
+            return
+
+        keep = ~particles
+        self._data = self._data.iloc[keep]
+
+    @property
+    def size(self):
+        return len(self._data)
+
+    def __getitem__(self, item):
+        return self._data[item].values
+
+    def __setitem__(self, item, value):
+        self._data[item] = value
+
+    def __getattr__(self, item):
+        if item not in self:
+            raise AttributeError(f'Attribute not defined: {item}')
+        return self[item]
+
+    def __contains__(self, item):
+        return item in self._data
+
+    @property
+    def dt(self):
+        """Backwards-compatibility function for returning model.solver.step"""
+        return self.model.solver.step
+
+    @property
+    def timestamp(self):
+        """Backwards-compatibility function for returning solver time as numpy datetime"""
+        return np.int64(self.model.solver.time).astype('datetime64[s]')
+
+    @property
+    def timestep(self):
+        """Backwards-compatibility function for returning solver time as timestep"""
+        elapsed = self.model.solver.time - self.model.solver.start
+        return elapsed // self.model.solver.step