Add low rank spline implementation

Author: plainerman

configs/toy/mueller_single_gaussian_low_rank_spline.yaml

@@ -0,0 +1,22 @@
+save_dir: ./out/toy/mueller_single_gaussian_low_rank_spline
+
+test_system: mueller_brown
+ode: first_order
+parameterization: low_rank
+T: 275e-4
+xi: 5.0
+gamma: 1.0
+
+model: spline
+num_points: 25
+
+num_gaussians: 1
+trainable_weights: False
+base_sigma: 2.5e-2
+
+epochs: 100000
+save_interval: 10000
+BS: 512
+
+num_paths: 1000
+dt: 1e-4

main.py

@@ -46,6 +46,12 @@
                     help="Whether the weights of the mixture model are trainable.")
 
 # model parameters
+parser.add_argument('--model', type=str, choices=['mlp', 'spline'], default='mlp',
+                    help="The model that will be used. Note that spline will not work with all configurations.")
+
+parser.add_argument('--num_points', type=int, default=100, help="Number of points in the spline model.")
+
+# MLP arguments
 parser.add_argument('--hidden_layers', nargs='+', type=int, help='The dimensions of the hidden layer of the MLP.',
                     default=[128, 128, 128])
 parser.add_argument('--activation', type=str, default='swish', choices=['tanh', 'relu', 'swish'],
@@ -75,7 +81,7 @@
 # plotting
 parser.add_argument('--log_plots', type=str2bool, nargs='?', const=True, default=False,
                     help="Save plots in log scale where possible")
-
+parser.add_argument('--extension', type=str, default='pdf', help="Extension of the saved plots.")
 
 def main():
     print("!!!!Next todos: plot ALDP")
@@ -120,7 +126,10 @@ def main():
 
     # You can play around with any model here
     # The chosen setup will append a final layer so that the output is mu, sigma, and weights
-    model = MLP(args.hidden_layers, args.activation, args.resnet)
+    model = None
+    if args.model == 'MLP':
+        model = MLP(args.hidden_layers, args.activation, args.resnet)
+
     setup = qsetup.construct(system, model, xi, A, B, args)
 
     key = jax.random.PRNGKey(args.seed)
@@ -166,7 +175,7 @@ def main():
         print("Warning: Loss contains NaNs")
     plt.plot(ckpt['losses'])
     log_scale(args.log_plots, False, True)
-    show_or_save_fig(args.save_dir, 'loss_plot.pdf')
+    show_or_save_fig(args.save_dir, 'loss_plot', args.extension)
 
     print("!!!TODO: how to plot this nicely?")
     # t = args.T * jnp.linspace(0, 1, args.BS, dtype=jnp.float32).reshape((-1, 1))
@@ -192,14 +201,14 @@ def main():
     if system.plot:
         # In case we have a second order integration scheme, we remove the velocity for plotting
         system.plot(title='Deterministic Paths', trajectories=x_t_det[:, :, :system.A.shape[0]])
-        show_or_save_fig(args.save_dir, 'paths_deterministic.pdf')
+        show_or_save_fig(args.save_dir, 'paths_deterministic', args.extension)
 
     key, path_key = jax.random.split(key)
     x_t_stoch = setup.sample_paths(state_q, x_0, args.dt, args.T, args.BS, path_key)
 
     if system.plot:
         system.plot(title='Stochastic Paths', trajectories=x_t_stoch[:, :, :system.A.shape[0]])
-        show_or_save_fig(args.save_dir, 'paths_stochastic.pdf')
+        show_or_save_fig(args.save_dir, 'paths_stochastic', args.extension)
 
 
 if __name__ == '__main__':

model/__init__.py

@@ -1,5 +1,4 @@
 from typing import Tuple
-
 from flax import linen as nn
 
 

training/qsetup.py

@@ -73,7 +73,7 @@ def u_t(self, state_q: TrainState, t: ArrayLike, x_t: ArrayLike, deterministic:
         raise NotImplementedError
 
 
-def construct(system: System, model: nn.module, xi: float, A: ArrayLike, B: ArrayLike,
+def construct(system: System, model: Optional[nn.module], xi: float, A: ArrayLike, B: ArrayLike,
               args: argparse.Namespace) -> QSetup:
     from training.setups import diagonal, lowrank
 
@@ -83,14 +83,21 @@ def construct(system: System, model: nn.module, xi: float, A: ArrayLike, B: Arra
         transform = aldp.InternalCoordinateWrapper(system.A.reshape(1, -1))
 
     if args.parameterization == 'diagonal':
-        wrapped_module = diagonal.DiagonalWrapper(
+        if args.model == 'spline':
+            raise ValueError("Spline model is not supported with diagonal parameterization")
+        model = diagonal.DiagonalWrapper(
             model, args.T, transform, A, B, args.num_gaussians, args.trainable_weights, args.base_sigma
         )
-        return diagonal.DiagonalSetup(system, wrapped_module, xi, args.ode, args.T)
+        return diagonal.DiagonalSetup(system, model, xi, args.ode, args.T)
     elif args.parameterization == 'low_rank':
-        wrapped_module = lowrank.LowRankWrapper(
-            model, args.T, transform, A, B, args.num_gaussians, args.trainable_weights, args.base_sigma
-        )
-        return lowrank.LowRankSetup(system, wrapped_module, xi, args.ode, args.T)
+        if args.model == 'spline':
+            model = lowrank.LowRankSpline(
+                args.num_points, args.T, transform, A, B, args.num_gaussians, args.trainable_weights, args.base_sigma
+            )
+        else:
+            model = lowrank.LowRankWrapper(
+                model, args.T, transform, A, B, args.num_gaussians, args.trainable_weights, args.base_sigma
+            )
+        return lowrank.LowRankSetup(system, model, xi, args.ode, args.T)
     else:
         raise ValueError(f"Unknown parameterization: {args.parameterization}")

training/setups/diagonal.py

@@ -41,13 +41,7 @@ def _post_process(self, h: ArrayLike, t: ArrayLike):
         return mu, sigma, w_logits
 
 
-@dataclass
 class DiagonalSetup(DriftedSetup):
-    model_q: DiagonalWrapper
-
-    def __init__(self, system: System, model_q: DiagonalWrapper, xi: ArrayLike, ode: str, T: float):
-        super().__init__(system, model_q, xi, ode, T)
-
     def construct_loss(self, state_q: TrainState, gamma: float, BS: int) -> Callable[
         [Union[FrozenVariableDict, Dict[str, Any]], ArrayLike], ArrayLike]:
 

training/setups/lowrank.py

@@ -1,16 +1,64 @@
-from dataclasses import dataclass
 from jax.typing import ArrayLike
 from flax import linen as nn
 import jax.numpy as jnp
-from typing import Union, Dict, Any, Callable, Tuple
+from typing import Union, Dict, Any, Callable, Tuple, Optional
 from flax.training.train_state import TrainState
 import jax
 from flax.typing import FrozenVariableDict
 from model.utils import WrappedModule
-from systems import System
 from training.setups.drift import DriftedSetup
 from training.utils import forward_and_derivatives
 
+interp = jax.vmap(jnp.interp, in_axes=(None, None, 1))
+
+
+class LowRankSpline(nn.Module):
+    n_points: int
+    T: float
+    transform: Optional[Callable[[Any], Any]]
+    A: ArrayLike
+    B: ArrayLike
+    num_mixtures: int
+    trainable_weights: bool
+    base_sigma: float
+
+    @nn.compact
+    def __call__(self, t):
+        print("WARNING: Mixtures for low rank not yet implemented!")
+        assert self.num_mixtures == 1, "Mixtures for low rank not yet implemented!"
+
+        ndim = self.A.shape[0]
+        t = t / self.T
+        t_grid = jnp.linspace(0, 1, self.n_points, dtype=jnp.float32)
+        S_0 = jnp.log(self.base_sigma) * jnp.eye(ndim, dtype=jnp.float32)
+        S_0_vec = S_0[jnp.tril_indices(ndim)]
+        mu_params = self.param('mu_params', lambda rng: jnp.linspace(self.A, self.B, self.n_points)[1:-1])
+        S_params = self.param('S_params', lambda rng: jnp.linspace(S_0_vec, S_0_vec, self.n_points)[1:-1])
+        y_grid = jnp.concatenate([self.A.reshape(1, -1), mu_params, self.B.reshape(1, -1)])
+        S_grid = jnp.concatenate([S_0_vec[None, :], S_params, S_0_vec[None, :]])
+
+        @jax.vmap
+        def get_tril(v):
+            a = jnp.zeros((ndim, ndim), dtype=jnp.float32)
+            a = a.at[jnp.tril_indices(ndim)].set(v)
+            return a
+
+        mu = interp(t.flatten(), t_grid, y_grid).T
+        S = interp(t.flatten(), t_grid, S_grid).T
+        S = get_tril(S)
+        S = jnp.tril(2 * jax.nn.sigmoid(S) - 1.0, k=-1) + jnp.eye(ndim, dtype=jnp.float32)[None, ...] * jnp.exp(S)
+
+        if self.trainable_weights:
+            w_logits = self.param('w_logits', nn.initializers.zeros_init(), (self.num_mixtures,), dtype=jnp.float32)
+        else:
+            w_logits = jnp.zeros(self.num_mixtures, dtype=jnp.float32)
+
+        out = (mu, S, w_logits)
+        if self.transform:
+            out = self.transform(out)
+
+        return out
+
 
 class LowRankWrapper(WrappedModule):
     A: ArrayLike
@@ -57,19 +105,10 @@ def get_tril(v):
         else:
             w_logits = jnp.zeros(num_mixtures, dtype=jnp.float32)
 
-        print('mu.shape', mu.shape)
-        print('S.shape', S.shape)
-
         return mu, S, w_logits
 
 
-@dataclass
 class LowRankSetup(DriftedSetup):
-    model_q: LowRankWrapper
-
-    def __init__(self, system: System, model_q: LowRankWrapper, xi: ArrayLike, order: str, T: float):
-        super().__init__(system, model_q, xi, order, T)
-
     def construct_loss(self, state_q: TrainState, gamma: float, BS: int) -> Callable[
         [Union[FrozenVariableDict, Dict[str, Any]], ArrayLike], ArrayLike]:
         def loss_fn(params_q: Union[FrozenVariableDict, Dict[str, Any]], key: ArrayLike) -> ArrayLike:

utils/plot.py

@@ -14,9 +14,9 @@ def log_scale(log_plot: bool, x: bool, y: bool):
             plt.gca().set_yscale('log')
 
 
-def show_or_save_fig(save_dir: Optional[str], name: str):
+def show_or_save_fig(save_dir: Optional[str], name: str, extension: str):
     if save_dir is not None:
-        plt.savefig(f'{save_dir}/{name}', bbox_inches='tight')
+        plt.savefig(f'{save_dir}/{name}.{extension}', bbox_inches='tight')
         plt.clf()
     else:
         plt.show()