added support for neural odes

Dockerfile-pytorchlatest

@@ -0,0 +1,11 @@
+# Dockerfile-pytorch1.13.1
+FROM pytorch/pytorch:latest
+
+WORKDIR /app
+
+# Copy the source code and install the package
+COPY . /flowc
+RUN pip install /flowc
+
+CMD ["/bin/bash"]
+

README.md

@@ -99,8 +99,9 @@ log_prob = flow.log_prob(inputs)
 ```
 
 To sample from the flow:
+
 ```python
-samples = flow.sample(num_samples)
+samples = flow.sample_like(num_samples)
 ```
 
 Additional examples of the workflow are provided in [examples folder](examples/).

docker/Dockerfile-devel

@@ -0,0 +1,13 @@
+# Dockerfile-pytorch1.13.1
+FROM pytorch/pytorch:latest
+
+WORKDIR /app
+
+# Copy the source code and install the package
+COPY . /app
+RUN pip install pdoc3
+RUN pip install -e .
+
+CMD ["/bin/bash"]
+
+EXPOSE 8080

docker/Dockerfile-pytorchlatest

@@ -3,9 +3,13 @@ FROM pytorch/pytorch:latest
 
 WORKDIR /app
 
-# Copy the source code and install the package
+RUN apt update && apt install build-essential -y --no-install-recommends
+
+# Copy the source code and install the packagepdoc
 COPY . /flowc
 RUN pip install /flowc
 
 CMD ["/bin/bash"]
 
+EXPOSE 8080
+

docker/requirements.txt

@@ -0,0 +1,15 @@
+matplotlib
+numpy
+pandas
+torchdiffeq
+umnn
+tqdm
+ninja
+scikit-learn
+h5py
+torchtestcase
+parameterized
+testingflake8
+pytest
+pytestcov
+black

examples/conditional_toy_2d.py

@@ -133,7 +133,8 @@ def plot_model(flow: Flow, dataset: PlaneDataset):
     plt.title('iteration {}'.format(i + 1))
     plt.tight_layout()
     # plt.show()
-    plt.show()
+    plt.savefig(f"figures/conditional_{selected_data}.png")
+    plt.close()
 
 
 if __name__ == "__main__":

flowcon/flows/base.py

@@ -50,9 +50,9 @@ def _log_prob(self, inputs, context):
     def _sample(self, num_samples, context):
         embedded_context = self._embedding_net(context)
         if self._context_used_in_base:
-            noise = self._distribution.sample(num_samples, context=embedded_context)
+            noise = self._distribution.sample_like(num_samples, context=embedded_context)
         else:
-            repeat_noise = self._distribution.sample(num_samples * embedded_context.shape[0])
+            repeat_noise = self._distribution.sample_like(num_samples * embedded_context.shape[0])
             noise = torch.reshape(
                 repeat_noise,
                 (embedded_context.shape[0], -1, repeat_noise.shape[1])

flowcon/nn/__init__.py

@@ -1,5 +1,5 @@
 from flowcon.nn.nets import *
-
+from flowcon.nn.neural_odes import odefunc
 __all__ = ['DenseNet',
            'MixedConditionalDenseNet',
            'InputConditionalDenseNet',
@@ -17,5 +17,6 @@
            'ConvResidualNet',
            'ResidualNet',
            'MLP',
-           'FCBlock'
+           'FCBlock',
+
            ]

flowcon/nn/nets/invertible_densenet.py

@@ -8,6 +8,7 @@
 from abc import ABC, abstractmethod
 from pprint import pformat
 
+import flowcon.utils.torchutils
 from flowcon.nn.nets import activations
 from flowcon.nn.nets.extended_basic_nets import ExtendedSequential, ExtendedLinear
 from flowcon.nn.nets.spectral_norm import scaled_spectral_norm
@@ -392,7 +393,7 @@ def __init__(self, dimension, context_features,
 
     def forward(self, inputs, context=None):
         context = self.bn(context)
-        values_weights = self.dense_net(inputs).unsqueeze(-1)
+        values_weights = flowcon.utils.torchutils.unsqueeze(-1)
         weights = self.custom_attention.attention(context, values_weights)
         return weights
 
@@ -449,7 +450,7 @@ def forward(self, inputs, context=None):
         context = self.bn(context)
         context_embedding = self.context_embedding_net(context)
         concat_inputs = torch.cat([inputs, context_embedding], -1)
-        values_weights = self.dense_net(concat_inputs).unsqueeze(-1)
+        values_weights = flowcon.utils.torchutils.unsqueeze(-1)
         weights = self.custom_attention.attention(context, values_weights)
         return weights
 

flowcon/nn/neural_odes/__init__.py

@@ -0,0 +1,4 @@
+
+from flowcon.nn.neural_odes.wrappers.cnf_regularization import RegularizedODEfunc
+from flowcon.nn.neural_odes import diffeq_layers
+from flowcon.nn.neural_odes.odefunc import *

flowcon/CNF/neural_odes/odefunc.py → flowcon/nn/neural_odes/odefunc.py

@@ -2,10 +2,12 @@
 import numpy as np
 import torch
 import torch.nn as nn
-from flowcon.transforms import ActNorm
+from typing import *
+
+from flowcon.transforms.normalization import ActNorm
 
 from . import diffeq_layers
-from .squeeze import squeeze, unsqueeze
+from flowcon.utils.torchutils import unsqueeze, squeeze
 
 __all__ = ["ODEnet", "ODEfunc"]
 
@@ -28,13 +30,15 @@ class ODEnet(nn.Module):
     """
 
     def __init__(
-            self, hidden_dims, input_shape, strides, conv, layer_type="concat", nonlinearity="softplus", num_squeeze=0,
-            act_norm=False, scale_output=1
-    ):
+            self, hidden_dims, input_shape, strides=None, conv=False,
+            layer_type: Literal["ignore", "hyper",
+            "squash", "concat", "concat_v2", "concatsquash", "blend", "concatcoord"] = "concat",
+            nonlinearity: Literal["tanh", "relu", "softplus", "elu", "swish", "square", "identity"] = "softplus",
+            num_squeeze=0,
+            act_norm=False, scale_output=1):
         super(ODEnet, self).__init__()
         self.act_norm = act_norm
         if act_norm:
-            self.t_actnorm = ActNorm(1)
             self.x_actnorm = ActNorm(input_shape[0])
 
         self.scale_output = scale_output
@@ -97,8 +101,7 @@ def __init__(
 
     def forward(self, t, y):
         if self.act_norm:
-            t = self.t_actnorm(t.view(-1, 1))[0].view(t.shape)
-            y, _ = self.x_actnorm(y)
+            y, logabsdet_actnorm = self.x_actnorm(y)
         dx = y
         # squeeze
         for _ in range(self.num_squeeze):

flowcon/transforms/__init__.py

@@ -86,3 +86,5 @@
 from flowcon.transforms.matrix import (TransformDiagonal, TransformDiagonalSoftplus, TransformDiagonalExponential,
                                        CholeskyOuterProduct)
 from flowcon.transforms.lipschitz import (iResBlock)
+
+from flowcon.transforms.neuralode import NeuralODE