The outputs of `BaseModel.predict` and `BaseModel.modes` are unintuitive and not documented

[pietronvll]

Introduction

kooplearn implements different models to learn the linear operator $\mathsf{T}f = \mathbb{E}[f(X_{t + 1}) | X_{t}, \ldots, X_{t - \ell}]$, where $\ell$ is the lookback length. The transfer operator $\mathsf{T}$ can be powered to return $\mathsf{T}^{s}f = \mathbb{E}[f(X_{t + s}) | X_{t}, \ldots, X_{t - \ell}]$, which can be used to get predictions farther in the future.

Currently, the predict methods in kooplearn return a raw tensor, or a dictionary of tensors, if any observable has been passed at training time. See, for example,

kooplearn/kooplearn/models/dict_of_fns.py

Lines 225 to 278 in ca71864

	def predict(
	self,
	data: TensorContextDataset,
	t: int = 1,
	predict_observables: bool = True,
	reencode_every: int = 0,
	):
	"""
	Predicts the state or, if the system is stochastic, its expected value :math:`\mathbb{E}[X_t | X_0 = X]` after ``t`` instants given the initial conditions ``data.lookback(self.lookback_len)`` being the lookback slice of ``data``.
	If ``data.observables`` is not ``None``, returns the analogue quantity for the observable instead.
	Args:
	data (TensorContextDataset): Dataset of context windows. The lookback window of ``data`` will be used as the initial condition, see the note above.
	t (int): Number of steps in the future to predict (returns the last one).
	predict_observables (bool): Return the prediction for the observables in ``self.data_fit.observables``, if present. Default to ``True``.
	reencode_every (int): When ``t > 1``, periodically reencode the predictions as described in :footcite:t:`Fathi2023`. Only available when ``predict_observables = False``.
	Returns:
	The predicted (expected) state/observable at time :math:`t`. The result is composed of arrays with shape matching ``data.lookforward(self.lookback_len)`` or the contents of ``data.observables``. If ``predict_observables = True`` and ``data.observables != None``, the returned ``dict``will contain the special key ``__state__`` containing the prediction for the state as well.
	"""
	check_is_fitted(
	self, ["U", "cov_XY", "cov_X", "cov_Y", "data_fit", "lookback_len"]
	)
	observables = None
	if predict_observables and hasattr(self.data_fit, "observables"):
	observables = self.data_fit.observables
	parsed_obs, expected_shapes, X_inference, X_fit = parse_observables(
	observables, data, self.data_fit
	)
	phi_Xin = self.feature_map(X_inference)
	phi_X = self.feature_map(X_fit)
	results = {}
	for obs_name, obs in parsed_obs.items():
	if (reencode_every > 0) and (t > reencode_every):
	if (predict_observables is True) and (observables is not None):
	raise ValueError(
	"rencode_every only works when forecasting states, not observables. Consider setting predict_observables to False."
	)
	else:
	num_reencodings = floor(t / reencode_every)
	for k in range(num_reencodings):
	raise NotImplementedError
	else:
	obs_pred = primal.predict(t, self.U, self.cov_XY, phi_Xin, phi_X, obs)
	obs_pred = obs_pred.reshape(expected_shapes[obs_name])
	results[obs_name] = obs_pred
	if len(results) == 1:
	return results["__state__"]
	else:
	return results

Same happens with the modes:

kooplearn/kooplearn/models/dict_of_fns.py

Lines 329 to 370 in ca71864

	def modes(
	self,
	data: TensorContextDataset,
	predict_observables: bool = True,
	):
	"""
	Computes the mode decomposition of arbitrary observables of the Koopman/Transfer operator at the states defined by ``data``.
	Informally, if :math:`(\\lambda_i, \\xi_i, \\psi_i)_{i = 1}^{r}` are eigentriplets of the Koopman/Transfer operator, for any observable :math:`f` the i-th mode of :math:`f` at :math:`x` is defined as: :math:`\\lambda_i \\langle \\xi_i, f \\rangle \\psi_i(x)`. See :footcite:t:`Kostic2022` for more details.
	Args:
	data (TensorContextDataset): Dataset of context windows. The lookback window of ``data`` will be used as the initial condition, see the note above.
	predict_observables (bool): Return the prediction for the observables in ``self.data_fit.observables``, if present. Default to ``True``.
	Returns:
	(modes, eigenvalues): Modes and corresponding eigenvalues of the system at the states defined by ``data``. The result is composed of arrays with shape matching ``data.lookforward(self.lookback_len)`` or the contents of ``data.observables``. If ``predict_observables = True`` and ``data.observables != None``, the returned ``dict`` will contain the special key ``__state__`` containing the modes for the state as well.
	"""
	check_is_fitted(self, ["U", "cov_XY", "data_fit", "lookback_len"])
	observables = None
	if predict_observables and hasattr(self.data_fit, "observables"):
	observables = self.data_fit.observables
	parsed_obs, expected_shapes, X_inference, X_fit = parse_observables(
	observables, data, self.data_fit
	)
	phi_Xin = self.feature_map(X_inference)
	phi_X = self.feature_map(X_fit)
	_gamma, _eigs = primal.estimator_modes(self.U, self.cov_XY, phi_X, phi_Xin)
	results = {}
	for obs_name, obs in parsed_obs.items():
	expected_shape = (self.rank,) + expected_shapes[obs_name]
	res = np.tensordot(_gamma, obs, axes=1).reshape(
	expected_shape
	) # [rank, num_initial_conditions, ...]
	results[obs_name] = res
	if len(results) == 1:
	return results["__state__"], _eigs
	else:
	return results, _eigs

On the contrary, training data are wrapped into an appropriate (and much better documented)

kooplearn/kooplearn/abc.py

Line 275 in ca71864

class ContextWindowDataset(ContextWindow):

Issues

There are multiple issues with the current approach:

1. The flow of data is inconsistent. A typical pipeline is as follows:
   1. One starts with one or multiple raw trajectories in the form of long arrays of shape (samples, features) (for example coming from ´kooplearn.datasets´)
   2. The trajectories are wrapped into contexts via TrajectoryContextDataset or the (unfinished, see add function for generating TrajectoryContextDataset from a list of trajectories #12) MultiTrajectoryContextDataset
   3. These context windows are either further wrapped into a DataLoader to be used with the neural-network models of kooplearn, or directly fed into BaseModel.fit
   4. Once a model is fitted, one calls model.predict(data: TensorContextWindow) or model.modes(data: TensorContextWindow).
   5. A raw tensor or dictionary of tensors is returned, with no obvious link to the input data.
2. Observables are not well supported. Right now we just check if the data has an observables attribute which can be accessed as a dictionary.

   kooplearn/kooplearn/models/dict_of_fns.py

   Lines 250 to 251 in ca71864

   	if predict_observables and hasattr(self.data_fit, "observables"):
   	observables = self.data_fit.observables

   This is undocumented, and observables should be properly registered into the ContextWindowDataset at initialization.
3. The Koopman modes are a structured object, while we return a tensor leaving all the post-processing to the user (see, for example, the compute_mode_info function in the switching system example).

Proposed solution

This will be a big new code release and not just a hotfix, it will take some time. I propose the following plan:

• Before writing any new code: re-evaluate the data flow and the role of each class in a sample scenario in which we have non-trivial observables to predict and for which we want to compute the modes. The goal of this step is to define a list of interventions which:
  1. Maximize the ease of use of the code
  2. Do not overhaul the structure of kooplearn. For example, it is fine if we refactor and simplify the context window objects; not fine if we re-think the data paradigm from scratch.
  3. Allow us to re-think critically the role of each of the objects and methods currently defined in kooplearn.data and kooplearn.abc. If we find out that some of these objects only add unnecessary abstraction, we should remove them.
• Implement these changes, and test them on:
  • All of the kernel/dictionary-of-function methods
  • The autoencoder methods
  • The neural-network feature maps
• Finish and test the support for multi-trajectory from add function for generating TrajectoryContextDataset from a list of trajectories #12. The TensorContextDataset is a nice abstraction, but it is pretty rare that it gets instantiated directly. Indeed, people usually have trajectories, and we should have support for handling trajectories as "first class citizens".
• Starting from @GregoirePacreau's compute_mode_info, design a class holding Koopman modes which is easy to work with.
• Write and update tests in tests/
• Update the documentation