Small TRAM refactorings (#192)

Author: MaaikeG

deeptime/markov/msm/tram/_bindings/src/tram_module.cpp

@@ -2,7 +2,7 @@
 #include "deeptime/markov/msm/tram/tram.h"
 #include "deeptime/markov/msm/tram/connected_set.h"
 #include "deeptime/markov/msm/tram/trajectory_mapping.h"
-
+ 
 PYBIND11_MODULE(_tram_bindings, m) {
     using namespace pybind11::literals;
     using namespace deeptime::markov::tram;
@@ -12,8 +12,8 @@ PYBIND11_MODULE(_tram_bindings, m) {
         py::class_<TRAM<double>>(tramMod, "TRAM")
                 .def(py::init<std::size_t, std::size_t>(), "n_therm_states"_a, "n_markov_states"_a)
                 .def(py::init<deeptime::np_array_nfc<double> &,
-                                deeptime::np_array_nfc<double> &, deeptime::np_array_nfc<double> &>(),
-                                "biased_conf_energies"_a, "lagrangian_mult_log"_a, "modified_state_counts_log"_a)
+                             deeptime::np_array_nfc<double> &, deeptime::np_array_nfc<double> &>(),
+                     "biased_conf_energies"_a, "lagrangian_mult_log"_a, "modified_state_counts_log"_a)
                 .def("estimate", &TRAM<double>::estimate,
                      "input"_a, "max_iter"_a = 1000, "max_err"_a = 1e-8, "callback_interval"_a = 1,
                      "track_log_likelihoods"_a = false, "callback"_a = nullptr)
@@ -22,17 +22,15 @@ PYBIND11_MODULE(_tram_bindings, m) {
                 .def_property_readonly("modified_state_counts_log", &TRAM<double>::modifiedStateCountsLog)
                 .def_property_readonly("lagrangian_mult_log", &TRAM<double>::lagrangianMultLog)
                 .def_property_readonly("therm_state_energies", &TRAM<double>::thermStateEnergies)
-                .def_property_readonly("markov_state_energies", &TRAM<double>::markovStateEnergies)
-                .def("compute_log_likelihood", &TRAM<double>::computeLogLikelihood, py::call_guard<py::gil_scoped_release>());
-
+                .def_property_readonly("markov_state_energies", &TRAM<double>::markovStateEnergies);
 
         py::class_<TRAMInput<double>, std::shared_ptr<TRAMInput<double>>>(tramMod, "TRAMInput").def(
-                py::init<deeptime::np_array_nfc<int> &&, deeptime::np_array_nfc<int> &&, DTrajs, BiasMatrices<double>>(),
-                "state_counts"_a, "transition_counts"_a, "dtrajs"_a, "bias_matrices"_a);
+                py::init<deeptime::np_array_nfc<int> &&, deeptime::np_array_nfc<int> &&, DTraj, BiasMatrix<double>>(),
+                "state_counts"_a, "transition_counts"_a, "dtraj"_a, "bias_matrix"_a);
 
-        tramMod.def("compute_sample_weights", &computeSampleWeights<double>, py::call_guard<py::gil_scoped_release>(),
-                    "therm_state_index"_a = -1, "dtrajs"_a, "bias_matrices"_a, "therm_state_energies"_a,
-                    "modified_state_counts_log"_a);
+        tramMod.def("compute_sample_weights_log", &computeSampleWeightsLog<double>, py::call_guard<py::gil_scoped_release>(),
+                    "dtraj"_a, "bias_matrix"_a, "therm_state_energies"_a,
+                    "modified_state_counts_log"_a, "therm_state_index"_a = -1);
 
         tramMod.def("find_state_transitions_post_hoc_RE",
                     &findStateTransitions<double, OverlapPostHocReplicaExchange<double>>,
@@ -46,5 +44,9 @@ PYBIND11_MODULE(_tram_bindings, m) {
                     "connectivity_factor"_a, "callback"_a);
 
         tramMod.def("find_trajectory_fragment_indices", &findTrajectoryFragmentIndices, "ttrajs"_a, "n_therm_states"_a);
+
+        tramMod.def("compute_log_likelihood", computeLogLikelihood<double>, py::call_guard<py::gil_scoped_release>(),
+                    "dtraj"_a, "biasMatrix"_a, "biasedConfEnergies"_a, "modifiedStateCountsLog"_a,
+                    "thermStateEnergies"_a, "stateCounts"_a, "transitionCounts"_a, "transitionMatrices"_a);
     }
 }

deeptime/markov/msm/tram/_tram.py

@@ -119,31 +119,6 @@ def __init__(
         self.log_likelihoods = []
         self.increments = []
 
-    @property
-    def compute_log_likelihood(self) -> Optional[float]:
-        r"""The parameter-dependent part of the TRAM likelihood.
-
-        The definition can be found in :footcite:`wu2016multiensemble`, Equation (9).
-
-        Returns
-        -------
-        log_likelihood : float
-            The parameter-dependent part of the log-likelihood.
-
-
-        Notes
-        -----
-        Parameter-dependent, i.e., the factor
-
-        .. math:: \prod_{x \in X} e^{-b^{k(x)}(x)}
-
-        does not occur in the log-likelihood as it is constant with respect to the parameters, leading to
-
-        .. math:: \log \prod_{k=1}^K \left(\prod_{i,j} (p_{ij}^k)^{c_{ij}^k}\right) \left(\prod_{i} \prod_{x \in X_i^k} \mu(x) e^{f_i^k} \right)
-        """
-        if self._tram_estimator is not None:
-            return self._tram_estimator.compute_log_likelihood()
-
     def fetch_model(self) -> Optional[TRAMModel]:
         r"""Yields the most recent :class:`MarkovStateModelCollection` that was estimated.
         Can be None if fit was not called.

deeptime/markov/msm/tram/_tram_dataset.py

@@ -4,9 +4,7 @@
 
 from deeptime.util import types, callbacks
 from deeptime.util.decorators import cached_property
-
 from deeptime.markov import TransitionCountEstimator, TransitionCountModel
-
 from ._tram_bindings import tram
 
 
@@ -21,6 +19,30 @@ def _determine_n_therm_states(dtrajs, ttrajs):
         return _determine_n_states(ttrajs)
 
 
+def transition_counts_from_count_models(n_therm_states, n_markov_states, count_models):
+    transition_counts = np.zeros((n_therm_states, n_markov_states, n_markov_states), dtype=np.int32)
+
+    for k in range(n_therm_states):
+        model_k = count_models[k]
+        if model_k.count_matrix.sum() > 0:
+            i_s, j_s = np.meshgrid(model_k.state_symbols, model_k.state_symbols, indexing='ij')
+            # place submodel counts in our full-sized count matrices
+            transition_counts[k, i_s, j_s] = model_k.count_matrix
+
+    return transition_counts
+
+
+def state_counts_from_count_models(n_therm_states, n_markov_states, count_models):
+    state_counts = np.zeros((n_therm_states, n_markov_states), dtype=np.int32)
+
+    for k in range(n_therm_states):
+        model_k = count_models[k]
+        if model_k.count_matrix.sum() > 0:
+            state_counts[k, model_k.state_symbols] = model_k.state_histogram
+
+    return state_counts
+
+
 def to_zero_padded_array(arrays, desired_shape):
     """Pad a list of numpy arrays with zeros to desired shape. Desired shape should be at least the size of the
     largest np array in the list.
@@ -55,8 +77,8 @@ def _invalidate_caches():
 class TRAMDataset:
     r""" Dataset for organizing data and obtaining properties from data that are needed for TRAM.
     The minimum required parameters for constructing a TRAMDataset are the `dtrajs` and `bias_matrices`. In this case,
-    `ttrajs` are inferred from the shape of the `dtrajs`, by assuming each trajectory in `dtrajs` corresponds to a unique
-    thermodynamic state, with the index corresponding to the index of occurrence in `dtrajs`.
+    `ttrajs` are inferred from the shape of the `dtrajs`, by assuming each trajectory in `dtrajs` corresponds to a
+    unique thermodynamic state, with the index corresponding to the index of occurrence in `dtrajs`.
 
     The values at identical indices in `dtrajs`, `ttrajs` and `bias_matrices` correspond to the sample. For example, at
     indices `(i, n)` we find information about the :math:`n`-th sample in trajectory :math:`i`. `dtrajs[i][n]` gives us
@@ -141,8 +163,11 @@ def __init__(self, dtrajs, bias_matrices, ttrajs=None, n_therm_states=None, n_ma
 
     @property
     def tram_input(self):
-        r""" The TRAMInput object containing the data needed for estimation. """
-        return tram.TRAMInput(self.state_counts, self.transition_counts, self.dtrajs, self.bias_matrices)
+        r""" The TRAMInput object containing the data needed for estimation.
+        For estimation purposes, it does not matter which thermodynamic state each sample was sampled at. The dtrajs and
+        bias_matrices are therefore flattened along the first dimension, to speed up estimation. """
+        return tram.TRAMInput(self.state_counts, self.transition_counts,
+                              np.concatenate(self.dtrajs), np.concatenate(self.bias_matrices))
 
     @property
     def n_therm_states(self):
@@ -169,16 +194,7 @@ def transition_counts(self):
         :getter: the transition counts
         :type: ndarray(n, m, m)
         """
-        transition_counts = np.zeros((self.n_therm_states, self.n_markov_states, self.n_markov_states), dtype=np.int32)
-
-        for k in range(self.n_therm_states):
-            model_k = self.count_models[k]
-            if model_k.count_matrix.sum() > 0:
-                i_s, j_s = np.meshgrid(model_k.state_symbols, model_k.state_symbols)
-                # place submodel counts in our full-sized count matrices
-                transition_counts[k, i_s, j_s] = model_k.count_matrix.T
-
-        return transition_counts
+        return transition_counts_from_count_models(self.n_therm_states, self.n_markov_states, self.count_models)
 
     @cached_property
     def state_counts(self):
@@ -192,14 +208,7 @@ def state_counts(self):
         matrices that are all the same shape, which is easier to handle (matrices are padded with zeros for all empty
         states that got dropped by the TransitionCountModels).
         """
-        state_counts = np.zeros((self.n_therm_states, self.n_markov_states), dtype=np.int32)
-
-        for k in range(self.n_therm_states):
-            model_k = self.count_models[k]
-            if model_k.count_matrix.sum() > 0:
-                state_counts[k, model_k.state_symbols] = model_k.state_histogram
-
-        return state_counts
+        return state_counts_from_count_models(self.n_therm_states, self.n_markov_states, self.count_models)
 
     def check_against_model(self, model):
         r""" Check the number of thermodynamic states of the model against that of the dataset. The number of
@@ -385,7 +394,7 @@ def _find_largest_connected_set(self, connectivity, connectivity_factor, progres
             all_state_counts = np.asarray([estimator.fit_fetch(dtraj).state_histogram for dtraj in self.dtrajs],
                                           dtype=object)
             # pad with zero's so they are all the same size and easier for the cpp module to handle
-            all_state_counts = to_zero_padded_array(all_state_counts, self.n_markov_states)
+            all_state_counts = to_zero_padded_array(all_state_counts, self.n_markov_states).astype(np.int32)
 
             # get list of all possible transitions between thermodynamic states. A transition is only possible when two
             # thermodynamic states have an overlapping markov state. Whether the markov state overlaps depends on the

deeptime/markov/msm/tram/_tram_model.py

@@ -4,6 +4,7 @@
 from deeptime.numeric import logsumexp
 from deeptime.markov.msm import MarkovStateModelCollection
 
+from ._tram_dataset import transition_counts_from_count_models, state_counts_from_count_models
 from ._tram_bindings import tram
 
 
@@ -51,7 +52,7 @@ def __init__(self, count_models, transition_matrices,
                  lagrangian_mult_log,
                  modified_state_counts_log,
                  therm_state_energies=None,
-                 markov_state_energies=None,
+                 markov_state_energies=None
                  ):
         self.n_therm_states = biased_conf_energies.shape[0]
         self.n_markov_states = biased_conf_energies.shape[1]
@@ -66,6 +67,9 @@ def __init__(self, count_models, transition_matrices,
         else:
             self._therm_state_energies = therm_state_energies
 
+        self._transition_matrices = transition_matrices
+        self._count_models = count_models
+
         self._msm_collection = self._construct_msm_collection(
             count_models, transition_matrices)
 
@@ -145,8 +149,15 @@ def compute_sample_weights(self, dtrajs, bias_matrices, therm_state=-1):
 
         .. math:: \mu(x) = \left( \sum_k R^k_{i(x)} \mathrm{exp}[f^k_{i(k)}-b^k(x)] \right)^{-1}
         """
-        return tram.compute_sample_weights(therm_state, dtrajs, bias_matrices, self._therm_state_energies,
-                                           self._modified_state_counts_log)
+        # flatten input data
+        dtraj = np.concatenate(dtrajs)
+        bias_matrix = np.concatenate(bias_matrices)
+
+        sample_weights = self._compute_sample_weights(dtraj, bias_matrix, therm_state)
+
+        # return in the original list shape
+        traj_start_stops = np.concatenate(([0], np.cumsum([len(traj) for traj in dtrajs])))
+        return [sample_weights[traj_start_stops[i - 1]:traj_start_stops[i]] for i in range(1, len(traj_start_stops))]
 
     def compute_observable(self, observable_values, dtrajs, bias_matrices, therm_state=-1):
         r""" Compute an observable value.
@@ -169,11 +180,11 @@ def compute_observable(self, observable_values, dtrajs, bias_matrices, therm_sta
             The index of the thermodynamic state in which the observable need to be computed. If `therm_state=-1`, the
             observable is computed for the unbiased (reference) state.
         """
-        sample_weights = self.compute_sample_weights(dtrajs, bias_matrices, therm_state)
+        # flatten input data
+        observable_values = np.concatenate(observable_values)
 
-        # flatten both
-        sample_weights = np.reshape(sample_weights, -1)
-        observable_values = np.reshape(observable_values, -1)
+        sample_weights = self._compute_sample_weights(np.concatenate(dtrajs), np.concatenate(bias_matrices),
+                                                      therm_state)
 
         return np.dot(sample_weights, observable_values)
 
@@ -200,20 +211,68 @@ def compute_PMF(self, dtrajs, bias_matrices, bin_indices, therm_state=-1):
             computed for the unbiased (reference) state.
         """
         # TODO: account for variable bin widths
-        sample_weights = np.reshape(self.compute_sample_weights(dtrajs, bias_matrices, therm_state), -1)
-        binned_samples = np.reshape(bin_indices, -1)
+        sample_weights = self._compute_sample_weights(np.concatenate(dtrajs), np.concatenate(bias_matrices),
+                                                      therm_state)
+
+        binned_samples = np.concatenate(bin_indices)
 
         n_bins = binned_samples.max() + 1
         pmf = np.zeros(n_bins)
 
         for i in range(len(pmf)):
             indices = np.where(binned_samples == i)
-            pmf[i] = -np.log(np.sum(sample_weights[indices]))
+            if len(indices[0]) > 0:
+                pmf[i] = -np.log(np.sum(sample_weights[indices]))
 
         # shift minimum to zero
         pmf -= pmf.min()
         return pmf
 
+    def compute_log_likelihood(self, dtrajs, bias_matrices):
+        r"""The (parameter-dependent part of the) likelihood to observe the given data.
+
+        The definition can be found in :footcite:`wu2016multiensemble`, Equation (9).
+
+        Parameters
+        ----------
+        dtrajs : list(np.ndarray)
+            The list of discrete trajectories. `dtrajs[i][n]` contains the Markov state index of the :math:`n`-th sample
+            in the :math:`i`-th trajectory.
+        bias_matrices : list(np.ndarray)
+            The bias energy matrices. `bias_matrices[i][n, k]` contains the bias energy of the :math:`n`-th sample from
+            the :math:`i`-th trajectory, evaluated at thermodynamic state :math:`k`, :math:`b^k(x_{i,n})`. The bias
+            energy matrices should have the same size as `dtrajs` in both the first and second dimension. The third
+            dimension is of size `n_therm_state`, i.e. for each sample, the bias energy in every thermodynamic state is
+            calculated and stored in the `bias_matrices`.
+
+        Returns
+        -------
+        log_likelihood : float
+            The parameter-dependent part of the log-likelihood.
+
+
+        Notes
+        -----
+        Parameter-dependent, i.e., the factor
+
+        .. math:: \prod_{x \in X} e^{-b^{k(x)}(x)}
+
+        does not occur in the log-likelihood as it is constant with respect to the parameters, leading to
+
+        .. math:: \log \prod_{k=1}^K \left(\prod_{i,j} (p_{ij}^k)^{c_{ij}^k}\right) \left(\prod_{i} \prod_{x \in X_i^k} \mu(x) e^{f_i^k} \right)
+        """
+        dtraj = np.concatenate(dtrajs)
+        bias_matrix = np.concatenate(bias_matrices)
+
+        transition_counts = transition_counts_from_count_models(self.n_therm_states, self.n_markov_states,
+                                                                self._count_models)
+
+        state_counts = state_counts_from_count_models(self.n_therm_states, self.n_markov_states, self._count_models)
+
+        return tram.compute_log_likelihood(dtraj, bias_matrix, self._biased_conf_energies,
+                                           self._modified_state_counts_log, self._therm_state_energies, state_counts,
+                                           transition_counts, self._transition_matrices)
+
     def _construct_msm_collection(self, count_models, transition_matrices):
         r""" Construct a MarkovStateModelCollection from the transition matrices and energy estimates.
         For each of the thermodynamic states, one MarkovStateModel is added to the MarkovStateModelCollection. The
@@ -237,3 +296,8 @@ def _construct_msm_collection(self, count_models, transition_matrices):
         return MarkovStateModelCollection(transition_matrices_connected, stationary_distributions,
                                           reversible=True, count_models=count_models,
                                           transition_matrix_tolerance=1e-8)
+
+    def _compute_sample_weights(self, dtraj, bias_matrix, therm_state=-1):
+        sample_weights = tram.compute_sample_weights_log(dtraj, bias_matrix, self._therm_state_energies,
+                                                     self._modified_state_counts_log, therm_state)
+        return np.exp(np.asarray(sample_weights))