official-source/DLWP-CS/DLWP/model/extensions.py

#
# Copyright (c) 2019 Jonathan Weyn <jweyn@uw.edu>
#
# See the file LICENSE for your rights.
#

"""
Extension classes for doing more with models, generators, and so on.
"""

import numpy as np
import xarray as xr
import pandas as pd
from .models import DLWPNeuralNet, DLWPFunctional
from .models_torch import DLWPTorchNN
from .generators import DataGenerator, SeriesDataGenerator, ArrayDataGenerator
from ..util import insolation
import warnings


class TimeSeriesEstimator(object):
    """
    Sophisticated wrapper class for producing time series forecasts from a DLWP model, using a Generator with
    metadata. This class allows predictions with non-matching inputs and outputs, including in the variable, level, and
    time_step dimensions.
    """

    def __init__(self, model, generator):
        """
        Initialize a TimeSeriesEstimator from a model and generator.

        :param model: DLWP model instance
        :param generator: DLWP DataGenerator instance
        """
        if not isinstance(model, (DLWPNeuralNet, DLWPFunctional, DLWPTorchNN)):
            raise TypeError("'model' must be a valid instance of a DLWP model class")
        if not isinstance(generator, (DataGenerator, SeriesDataGenerator, ArrayDataGenerator)):
            raise TypeError("'generator' must be a valid instance of a DLWP generator class")
        if isinstance(model, DLWPFunctional):
            warnings.warn('DLWPFunctional models are only partially supported by TimeSeriesEstimator. The '
                          'inputs/outputs to the model must be the same if the model predicts a sequence.')
        self.model = model
        self.generator = generator
        self._add_insolation = generator._add_insolation if hasattr(generator, '_add_insolation') else False
        self._uses_varlev = 'varlev' in generator.ds.dims
        self._has_targets = 'targets' in generator.ds.variables
        self._is_series = isinstance(generator, (SeriesDataGenerator, ArrayDataGenerator))
        self._output_sel = {}
        self._input_sel = {}
        self._dt = self.generator.ds['sample'][1] - self.generator.ds['sample'][0]
        if hasattr(self.generator, '_interval'):
            self._interval = self.generator._interval
        else:
            self._interval = 1
        if hasattr(self.generator, 'rank'):
            self.rank = self.generator.rank
        else:
            self.rank = 2

        # Generate the selections needed for inputs and outputs
        if self._uses_varlev:
            # Find the outputs we keep
            if self._is_series:  # a SeriesDataGenerator has user-specified I/O
                if not self.generator._output_sel:  # selection was empty
                    self._output_sel = {'varlev': np.array(self.generator.ds.coords['varlev'][:])}
                else:
                    self._output_sel = {k: np.array(v) for k, v in self.generator._output_sel.items()}
            else:
                if self._has_targets:
                    self._output_sel = {'varlev': np.array(self.generator.ds.targets.coords['varlev'][:])}
                else:
                    self._output_sel = {'varlev': np.array(self.generator.ds.coords['varlev'][:])}
            # Find the inputs we need
            if self._is_series:
                if not self.generator._input_sel:  # selection was empty
                    self._input_sel = {'varlev': np.array(self.generator.ds.coords['varlev'][:])}
                else:
                    self._input_sel = {k: np.array(v) for k, v in self.generator._input_sel.items()}
            else:
                self._input_sel = {'varlev': np.array(self.generator.ds.predictors.coords['varlev'][:])}
            # Find outputs that need to replace inputs
            self._outputs_in_inputs = {
                'varlev': np.array([v for v in self._output_sel['varlev'] if v in self._input_sel['varlev']])
            }
        else:  # Uses variable/level coordinates
            # Find the outputs we keep
            if self._is_series:  # a SeriesDataGenerator has user-specified I/O
                if not self.generator._output_sel:  # selection was empty
                    self._output_sel = {'variable': np.array(self.generator.ds.coords['variable'][:]),
                                        'level': np.array(self.generator.ds.coords['level'][:])}
                else:
                    self._output_sel = {k: np.array(v) for k, v in self.generator._output_sel.items()}
                    if 'variable' not in self._output_sel.keys():
                        self._output_sel['variable'] = np.array(self.generator.ds.coords['variable'][:])
                    if 'level' not in self._output_sel.keys():
                        self._output_sel['level'] = np.array(self.generator.ds.coords['level'][:])
            else:
                if self._has_targets:
                    self._output_sel = {'variable': np.array(self.generator.ds.targets.coords['variable'][:]),
                                        'level': np.array(self.generator.ds.targets.coords['level'][:])}
                else:
                    self._output_sel = {'variable': np.array(self.generator.ds.coords['variable'][:]),
                                        'level': np.array(self.generator.ds.coords['level'][:])}
            # Find the inputs we need
            if self._is_series:
                if not self.generator._input_sel:  # selection was empty
                    self._input_sel = {'variable': np.array(self.generator.ds.coords['variable'][:]),
                                       'level': np.array(self.generator.ds.coords['level'][:])}
                else:
                    self._input_sel = {k: np.array(v) for k, v in self.generator._input_sel.items()}
                    if 'variable' not in self._input_sel.keys():
                        self._input_sel['variable'] = np.array(self.generator.ds.coords['variable'][:])
                    if 'level' not in self._input_sel.keys():
                        self._input_sel['level'] = np.array(self.generator.ds.coords['level'][:])
            else:
                self._input_sel = {'variable': np.array(self.generator.ds.predictors.coords['variable'][:]),
                                   'level': np.array(self.generator.ds.predictors.coords['level'][:])}
            # Flatten variable/level
            lev, var = np.meshgrid(self._input_sel['level'], self._input_sel['variable'])
            varlev = np.array(['/'.join([v, str(l)]) for v, l in zip(var.flatten(), lev.flatten())])
            self._input_sel['varlev'] = varlev
            lev, var = np.meshgrid(self._output_sel['level'], self._output_sel['variable'])
            varlev = np.array(['/'.join([v, str(l)]) for v, l in zip(var.flatten(), lev.flatten())])
            self._output_sel['varlev'] = varlev
            # Find outputs that need to replace inputs
            self._outputs_in_inputs = {
                'variable': np.array([v for v in self._output_sel['variable'] if v in self._input_sel['variable']]),
                'level': np.array([v for v in self._output_sel['level'] if v in self._input_sel['level']]),
                'varlev': np.array([v for v in self._output_sel['varlev'] if v in self._input_sel['varlev']])
            }
        if self._add_insolation:
            self._input_sel['varlev'] = np.concatenate([self._input_sel['varlev'], np.array(['SOL'])])

        # Time step dimension
        self._input_time_steps = (generator._input_time_steps if isinstance(generator, SeriesDataGenerator)
                                  else model.time_dim)
        self._output_time_steps = (generator._output_time_steps if isinstance(generator, SeriesDataGenerator)
                                   else model.time_dim)

        # Channels last option
        self.channels_last = hasattr(self.generator, 'channels_last') and self.generator.channels_last
        self._forward_transpose = (0, self.rank + 1) + tuple(range(1, 1 + self.rank)) + (-1,)
        self._backward_transpose = (0,) + tuple(range(2, 2 + self.rank)) + (1, -1,)

        # Constants in the generator
        if hasattr(self.generator, 'constants') and self.generator.constants is not None:
            if self.channels_last:
                self.constants = self.generator.constants.transpose(tuple(range(1, 1 + self.rank)) + (0,))
            else:
                self.constants = self.generator.constants
        else:
            self.constants = None

    @property
    def shape(self):
        return (self.generator._n_sample,) + self.generator.shape

    @property
    def convolution_shape(self):
        return (self.generator._n_sample,) + self.generator.convolution_shape

    def predict(self, steps, samples=(), impute=False, keep_time_dim=False, prefer_first_times=True,
                f_hour_timedelta_type=False, **kwargs):
        """
        Step forward the time series prediction from the model 'steps' times, feeding predictions back in as
        inputs. Predicts for all the data provided in the generator. If there are inputs which are not produced by
        the model outputs, we include the available inputs from the generator data and either reduce the number of
        predicted samples accordingly (remove those whose inputs cannot be satisfied) or run the model using the mean
        values of the inputs which cannot be satisfied. If there are fewer output time steps than input time steps,
        then we build a time series forecast intelligently using part of the predictors and part of the prediction at
        every step. Note only the SeriesDataGenerator supports variable inputs/outputs.

        :param steps: int: number of times to step forward
        :param samples: list of int: which samples in the generator to predict for. For all samples, pass an empty
            list or tuple. May cause unexpected behavior when the input and output data do not match.
        :param impute: bool: if True, use the mean state for missing inputs in the forward integration
        :param keep_time_dim: bool: if True, keep the time_step dimension instead of integrating it with forecsat_hour
            to produce a continuous time series
        :param prefer_first_times: bool: in the case where the prediction contains more time_steps than the input,
            use the first available predicted times to initialize the next step, otherwise use the last times. If the
            output time_steps is less than the input time_steps, we always use all of the output times.
        :param f_hour_timedelta_type: bool: if True, converts f_hour dimension into a timedelta type. May not always be
            compatible with netCDF applications.
        :param kwargs: passed to Keras.predict()
        :return: ndarray: predicted states with forecast_step as the first dimension
        """
        if int(steps) < 1:
            raise ValueError('must use positive integer for steps')

        # Effective forward time steps for each step
        if self._output_time_steps <= self._input_time_steps:
            keep_inputs = True
            es = self._output_time_steps
            in_times = np.arange(self._input_time_steps) - (self._input_time_steps - self._output_time_steps)
        else:
            keep_inputs = False
            if prefer_first_times:
                es = self._input_time_steps
                in_times = np.arange(self._input_time_steps)
            else:
                es = self._output_time_steps
                in_times = np.arange(self._input_time_steps) + (self._output_time_steps - self._input_time_steps)
        effective_steps = int(np.ceil(steps / es))

        # Load data from the generator, without any scaling as this will be done by the model's predict method
        predictors, t = self.generator.generate(samples, scale_and_impute=False)
        p = predictors[0] if isinstance(predictors, (list, tuple)) else predictors
        p_shape = tuple(p.shape)

        # Add metadata. The sample dimension denotes the *start* time of the sample, for insolation purposes. This is
        # then corrected when providing the final output time series.
        sample_coord = self.generator.ds.sample[:self.generator._n_sample] if len(samples) == 0 else \
            self.generator.ds.sample[samples]
        if not self._is_series:
            sample_coord = sample_coord - self._dt * (self._input_time_steps - 1)
        if self.channels_last:
            # Split time step/varlev at the end and transpose
            if not self.generator._keep_time_axis:
                p = p.reshape((p_shape[0],) + self.generator.convolution_shape[-self.rank-1:-1] +
                              (self._input_time_steps, -1)).transpose(self._forward_transpose)
            p_da = xr.DataArray(
                p,
                coords=([sample_coord, in_times] +
                        [np.arange(d) for d in self.generator.convolution_shape[-self.rank-1:-1]] +
                        [self._input_sel['varlev']]),
                dims=['sample', 'time_step'] + ['x%d' % d for d in range(self.rank)] + ['varlev']
            )
        else:
            p = p.reshape((p_shape[0], self._input_time_steps, -1,) + self.generator.convolution_shape[-self.rank:])
            p_da = xr.DataArray(
                p,
                coords=([sample_coord, in_times, self._input_sel['varlev']] +
                        [np.arange(d) for d in self.generator.convolution_shape[-self.rank:]]),
                dims=['sample', 'time_step', 'varlev'] + ['x%d' % d for d in range(self.rank)]
            )
        if self.rank == 2:
            p_da = p_da.rename({'x0': 'lat', 'x1': 'lon'}).assign_coords(lat=self.generator.ds.lat,
                                                                         lon=self.generator.ds.lon)

        # Calculate mean for imputing
        if impute:
            p_mean = p.mean(axis=0)

        # Target shape
        t_shape = t[0].shape if isinstance(t, (list, tuple)) else t.shape
        t = None

        # A DLWPFunctional model which does not have the same inputs/outputs must have some programmatic way of fixing
        # this issue built-in (if it is trained to minimize several iterations of the model). Thus we fall back to the
        # regular predict_timeseries method. However, such a model that does not predict a sequence is perfectly valid
        # for use here.
        if isinstance(self.model, DLWPFunctional) and self.model._n_steps > 1:
            if not self.generator._add_insolation:
                # For the DLWPFunctional model, just use its predict_timeseries API, which handles effective steps
                # TODO: correctly handle channels_last
                result = self.model.predict_timeseries(predictors, steps, keep_time_dim=True,
                                                       **kwargs).reshape((-1,) + t_shape)[:effective_steps, ...]
                result = result.transpose((0, 1) + tuple(range(2, len(result.shape)))).copy()
            else:
                # If insolation is requested, intelligently add it in the same way the generator does
                sequence_steps = int(np.ceil(steps / self.model._n_steps / self.model.time_dim))

                # Giant forecast array
                result = np.full((t_shape[0], sequence_steps, self.model._n_steps) + t_shape[1:],
                                 np.nan, dtype=np.float32)

                # Iterate
                new_t = p_da.sample[:]
                for s in range(sequence_steps):
                    if 'verbose' in kwargs and kwargs['verbose'] > 0:
                        print('Time step %d/%d' % (s + 1, sequence_steps))
                    result[:, s] = np.stack(self.model.predict(predictors, **kwargs), axis=1)

                    # Assign new insolation to list of inputs
                    new_t = new_t + self._output_time_steps * self.model._n_steps * self._dt
                    new_insolation = [np.concatenate(
                        [np.expand_dims(insolation(new_t + (n + m * self._input_time_steps) * self._dt,
                                                   self.generator.ds.lat.values,
                                                   self.generator.ds.lon.values)[:, None],
                                        axis=-1 if self.channels_last else 1)
                         for n in range(self._input_time_steps)],
                        axis=1) for m in range(self.model._n_steps)]

                    if self.channels_last:
                        if self.generator._keep_time_axis:
                            r = result[:, s, -1]
                            predictors = [np.concatenate([r, new_insolation[0]], axis=-1).reshape(
                                (-1,) + self.convolution_shape[1:])] + new_insolation[1:]
                        else:
                            r = result[:, s, -1].reshape(t_shape[:-1] + (self._output_time_steps, -1)).transpose(
                                self._forward_transpose
                            )
                            predictors = [np.concatenate([r, new_insolation[0]], axis=-1).transpose(
                                self._backward_transpose).reshape((-1,) + self.convolution_shape[1:])] \
                                + new_insolation[1:]
                    else:
                        predictors = [
                            np.concatenate([result[:, s, -1].reshape((-1,) + self.shape[1:]), new_insolation[0]],
                                           axis=2).reshape((-1,) + self.convolution_shape[1:])
                        ] + new_insolation[1:]

                    # Add constants
                    if self.constants is not None:
                        predictors.append(np.repeat(np.expand_dims(self.constants, axis=0), len(p_da.sample), axis=0))

                n_dim_1 = result.size // int(np.prod(t_shape))
                result.shape = (t_shape[0], n_dim_1,) + t_shape[1:]
                result = result[:, :effective_steps, ...]
        else:
            # Giant forecast array
            result = np.full((t_shape[0], effective_steps,) + t_shape[1:], np.nan, dtype=np.float32)

            # Iterate prediction forward for a regular DLWP Sequential NN
            for s in range(effective_steps):
                if 'verbose' in kwargs and kwargs['verbose'] > 0:
                    print('Time step %d/%d' % (s + 1, effective_steps))
                if self.channels_last and not self.generator._keep_time_axis:
                    result[:, s] = self.model.predict(p_da.values.transpose(self._backward_transpose).reshape(p_shape),
                                                      **kwargs)
                else:
                    result[:, s] = self.model.predict(p_da.values.reshape(p_shape), **kwargs)

                # Add metadata to the prediction
                if self.channels_last:
                    if not self.generator._keep_time_axis:
                        r = result[:, s].reshape((p_shape[0],) + self.generator.convolution_shape[-self.rank-1:-1] +
                                                 (self._output_time_steps, -1)).transpose(self._forward_transpose)
                    else:
                        r = result[:, s]
                    r_da = xr.DataArray(
                        r,
                        coords=([p_da.sample + (es + self._interval - 1) * self._dt,
                                 np.arange(self._output_time_steps)] +
                                [np.arange(d) for d in self.generator.convolution_shape[-self.rank-1:-1]] +
                                [self._output_sel['varlev']]),
                        dims=['sample', 'time_step'] + ['x%d' % d for d in range(self.rank)] + ['varlev']
                    )
                else:
                    r_da = xr.DataArray(
                        result[:, s].reshape((p_shape[0], self._output_time_steps, -1,) +
                                             self.generator.convolution_shape[-self.rank:]),
                        coords=([p_da.sample + (es + self._interval - 1) * self._dt,
                                 np.arange(self._output_time_steps), self._output_sel['varlev']] +
                                [np.arange(d) for d in self.generator.convolution_shape[-self.rank:]]),
                        dims=['sample', 'time_step', 'varlev'] + ['x%d' % d for d in range(self.rank)]
                    )
                if self.rank == 2:
                    r_da = r_da.rename({'x0': 'lat', 'x1': 'lon'}).assign_coords(lat=self.generator.ds.lat,
                                                                                 lon=self.generator.ds.lon)

                # Re-index the predictors to the new forward time step
                p_da = p_da.reindex(sample=r_da.sample, method=None)

                # Impute values extending beyond data availability
                if impute:
                    # Calculate mean values for the added time steps after re-indexing
                    p_da[-es:] = np.concatenate([p_mean[np.newaxis, ...]] * es)

                # Take care of the known insolation for added time steps
                if self._add_insolation:
                    p_da.loc[{'varlev': 'SOL'}][-es:] = \
                        np.concatenate([insolation(p_da.sample[-es:] + n * self._dt, self.generator.ds.lat.values,
                                                   self.generator.ds.lon.values)[:, np.newaxis]
                                        for n in range(self._input_time_steps)], axis=1)

                # Replace the predictors that exist in the result with the result. Any that do not exist are
                # automatically inherited from the known predictor data (or imputed data).
                if keep_inputs:
                    loc_dict = dict(varlev=self._outputs_in_inputs['varlev'], time_step=p_da.time_step[-es:])
                    p_da.loc[loc_dict] = r_da.loc[{'varlev': self._outputs_in_inputs['varlev']}]
                else:
                    if prefer_first_times:
                        p_da.loc[{'varlev': self._outputs_in_inputs['varlev']}] = \
                            r_da.loc[{'varlev': self._outputs_in_inputs['varlev']}][:, :self._input_time_steps]
                    else:
                        p_da.loc[{'varlev': self._outputs_in_inputs['varlev']}] = \
                            r_da.loc[{'varlev': self._outputs_in_inputs['varlev']}][:, -self._input_time_steps:]

        # Return a DataArray. Keep the actual model initialization, that is, the last available time in the inputs,
        # as the time
        rv = result.view()
        if self.channels_last:
            if not self.generator._keep_time_axis:
                rv.shape = (p_shape[0], effective_steps,) + \
                           self.generator.output_convolution_shape[-self.rank-1:-1] + (self._output_time_steps, -1,)
        else:
            rv.shape = (p_shape[0], effective_steps, self._output_time_steps, -1,) + \
                self.generator.output_convolution_shape[-self.rank:]
        if f_hour_timedelta_type:
            dt = self._dt.values
        else:
            dt = np.array(self._dt.values.astype('timedelta64[h]').astype('float'))
        if keep_time_dim:
            if self.channels_last:
                result_da = xr.DataArray(
                    rv.transpose((1, 0) + tuple(range(2, len(rv.shape)))) if self.generator._keep_time_axis else
                    rv.transpose((1, 0, -2) + tuple(range(2, 2 + self.rank)) + (-1,)),
                    coords=[
                               np.arange(dt, (effective_steps * (es + self._interval - 1) + 1) * dt,
                                         (es + self._interval - 1) * dt),
                               sample_coord + (self._input_time_steps - 1) * self._dt,
                               range(self._output_time_steps),
                           ]
                    + [np.arange(d) for d in self.generator.output_convolution_shape[-self.rank-1:-1]]
                    + [self._output_sel['varlev']],
                    dims=['f_hour', 'time', 'time_step'] + ['x%d' % d for d in range(self.rank)] + ['varlev'],
                    name='forecast'
                )
            else:
                result_da = xr.DataArray(
                    rv.transpose((1, 0) + tuple(range(2, len(rv.shape)))),
                    coords=[
                        np.arange(dt, (effective_steps * (es + self._interval - 1) + 1) * dt,
                                  (es + self._interval - 1) * dt),
                        sample_coord + (self._input_time_steps - 1) * self._dt,
                        range(self._output_time_steps),
                        self._output_sel['varlev'],
                    ] + [np.arange(d) for d in self.generator.output_convolution_shape[-self.rank:]],
                    dims=['f_hour', 'time', 'time_step', 'varlev'] + ['x%d' % d for d in range(self.rank)],
                    name='forecast'
                )
            if self.rank == 2:
                result_da = result_da.rename({'x0': 'lat', 'x1': 'lon'}).assign_coords(lat=self.generator.ds.lat,
                                                                                       lon=self.generator.ds.lon)
        else:
            # To create a correct time series, we must retain only the effective steps
            if not keep_inputs:
                if prefer_first_times:
                    rv = rv[:, :, :es]
            if self.channels_last:
                if self.generator._keep_time_axis:
                    rv.shape = (t_shape[0], rv.shape[1] * rv.shape[2]) + rv.shape[3:]
                result_da = xr.DataArray(
                    rv.transpose((1, 0) + tuple(range(2, len(rv.shape)))) if self.generator._keep_time_axis else
                    rv.transpose((1, -2, 0) + tuple(range(2, 2 + self.rank)) + (-1,)).reshape(
                        (rv.shape[1] * rv.shape[-2], rv.shape[0]) +
                        self.generator.output_convolution_shape[-self.rank-1:-1] + (-1,)
                    ),
                    coords=[
                               np.array([(np.arange(0, es) + self._interval + e * (es - 1 + self._interval)) * dt
                                         for e in range(effective_steps)]).flatten(),
                               sample_coord + (self._input_time_steps - 1) * self._dt,
                           ]
                    + [np.arange(d) for d in self.generator.output_convolution_shape[-self.rank-1:-1]]
                    + [self._output_sel['varlev']],
                    dims=['f_hour', 'time'] + ['x%d' % d for d in range(self.rank)] + ['varlev'],
                    name='forecast'
                )
            else:
                rv.shape = (t_shape[0], rv.shape[1] * rv.shape[2]) + rv.shape[3:]
                result_da = xr.DataArray(
                    rv.transpose((1, 0) + tuple(range(2, len(rv.shape)))),
                    coords=[
                        np.array([(np.arange(0, es) + self._interval + e * (es - 1 + self._interval)) * dt
                                  for e in range(effective_steps)]).flatten(),
                        sample_coord + (self._input_time_steps - 1) * self._dt,
                        self._output_sel['varlev'],
                    ] + [np.arange(d) for d in self.generator.output_convolution_shape[-self.rank:]],
                    dims=['f_hour', 'time', 'varlev'] + ['x%d' % d for d in range(self.rank)],
                    name='forecast'
                )
            if self.rank == 2:
                result_da = result_da.rename({'x0': 'lat', 'x1': 'lon'}).assign_coords(lat=self.generator.ds.lat,
                                                                                       lon=self.generator.ds.lon)
            result_da = result_da.isel(f_hour=slice(0, steps))

        # Expand back out to variable/level pairs
        if self._uses_varlev:
            return result_da
        else:
            var, lev = self._output_sel['variable'], self._output_sel['level']
            vl = pd.MultiIndex.from_product((var, lev), names=('variable', 'level'))
            result_da = result_da.assign_coords(varlev=vl).unstack('varlev')
            spatial_dims = [d for d in result_da.dims if d not in ['f_hour', 'time', 'variable', 'level']]
            if self.channels_last:
                transpose_dims = ('f_hour', 'time') + tuple(spatial_dims) + ('variable', 'level')
                result_da = result_da.transpose('f_hour', 'time', 'lat', 'lon', 'variable', 'level')
            else:
                transpose_dims = ('f_hour', 'time') + ('variable', 'level') + tuple(spatial_dims)
            result_da = result_da.transpose(*transpose_dims)
            return result_da


class SeriesDataGeneratorWithInference(SeriesDataGenerator):
    """
    An extension of the SeriesDataGenerator that couples a second model for inference of part of the sequence target
    data. For example, with a model that predicts a sequence, a fixed model with fixed weights can be used to predict
    the first step of a sequence while a new model is trained on the existing model first step and then real data for
    the following steps.
    """

    def __init__(self, inference_model, inference_steps, *args, **kwargs):
        """
        Initialize an SeriesDataGenerator with an inference model. The arguments and kwargs must be those passed to the
        ArrayDataGenerator. The parameter inference_steps governs which parts of the target sequence are replaced with
        the inference model prediction. Numbers in inference_steps should be integers starting with 0 that correspond
        to which steps in the target data should be replaced my the inference model prediction.

        :param inference_model: DLWP model
        :param inference_steps: iterable: steps in sequence to replace with inference
        :param args: passed to SeriesDataGenerator()
        :param kwargs: passed to SeriesDataGenerator()
        """
        if not isinstance(inference_model, (DLWPNeuralNet, DLWPFunctional, DLWPTorchNN)):
            raise TypeError("inference model should be a DLWP model")

        self.inference_model = inference_model
        super(SeriesDataGeneratorWithInference, self).__init__(*args, **kwargs)

        if self._sequence is None:
            raise ValueError("'SeriesDataGeneratorWithInference' is only usable if the generator produces a "
                             "sequence of targets")
        if np.any(np.array(inference_steps)) < 0 or np.any(np.array(inference_steps)) >= self._sequence:
            raise ValueError("got steps parameter (%s) outside of range 0 to %s" % (inference_steps, self._sequence))
        self.inference_steps = inference_steps

    def generate(self, samples, scale_and_impute=True):
        # Generate data normally
        p, t = super(SeriesDataGeneratorWithInference, self).generate(samples, scale_and_impute)

        # Make a prediction with the inference model
        predicted = self.inference_model.predict(p)

        # Insert inference prediction
        for s in self.steps:
            t[s] = predicted[s][:]

        # Return modified sample
        return p, t


class ArrayDataGeneratorWithInference(ArrayDataGenerator):
    """
    An extension of the ArrayDataGenerator that couples a second model for inference of part of the sequence target
    data. For example, with a model that predicts a sequence, a fixed model with fixed weights can be used to predict
    the first step of a sequence while a new model is trained on the existing model first step and then real data for
    the following steps.
    """

    def __init__(self, inference_model, inference_steps, *args, **kwargs):
        """
        Initialize an ArrayDataGenerator with an inference model. The arguments and kwargs must be those passed to the
        ArrayDataGenerator. The parameter inference_steps governs which parts of the target sequence are replaced with
        the inference model prediction. Numbers in inference_steps should be integers starting with 0 that correspond
        to which steps in the target data should be replaced my the inference model prediction.

        :param inference_model: DLWP model
        :param inference_steps: iterable: steps in sequence to replace with inference
        :param args: passed to ArrayDataGenerator()
        :param kwargs: passed to ArrayDataGenerator()
        """
        if not isinstance(inference_model, (DLWPNeuralNet, DLWPFunctional, DLWPTorchNN)):
            raise TypeError("inference model should be a DLWP model")

        self.inference_model = inference_model
        super(ArrayDataGeneratorWithInference, self).__init__(*args, **kwargs)

        if self._sequence is None:
            raise ValueError("'ArrayDataGeneratorWithInference' is only usable if the generator produces a "
                             "sequence of targets")
        if np.any(np.array(inference_steps)) < 0 or np.any(np.array(inference_steps)) >= self._sequence:
            raise ValueError("got steps parameter (%s) outside of range 0 to %s" % (inference_steps, self._sequence))
        self.inference_steps = inference_steps

    def generate(self, samples):
        # Generate data normally
        p, t = super(ArrayDataGeneratorWithInference, self).generate(samples)

        # Make a prediction with the inference model
        predicted = self.inference_model.predict(p)

        # Insert inference prediction
        for s in self.inference_steps:
            t[s] = predicted[s][:]

        # Return modified sample
        return p, t