update SDE class within notebooks and remove from utils

atong01 · Dec 14, 2023 · 980e3d9 · 980e3d9
1 parent 91fd396
commit 980e3d9
Show file tree

Hide file tree

Showing 5 changed files with 130 additions and 103 deletions.
diff --git a/examples/notebooks/SF2M_2D_example.ipynb b/examples/notebooks/SF2M_2D_example.ipynb
diff --git a/examples/notebooks/conditional_mnist.ipynb b/examples/notebooks/conditional_mnist.ipynb
diff --git a/examples/notebooks/mnist_example.ipynb b/examples/notebooks/mnist_example.ipynb
diff --git a/examples/notebooks/single-cell_example.ipynb b/examples/notebooks/single-cell_example.ipynb
diff --git a/torchcfm/utils.py b/torchcfm/utils.py
@@ -64,24 +64,3 @@ def plot_trajectories(traj):
     plt.xticks([])
     plt.yticks([])
     plt.show()
-
-
-class SDE(torch.nn.Module):
-    noise_type = "diagonal"
-    sde_type = "ito"
-
-    def __init__(self, ode_drift, score, input_size=(3, 32, 32), sigma=0.1):
-        super().__init__()
-        self.drift = ode_drift
-        self.score = score
-        self.input_size = input_size
-        self.sigma = sigma
-
-    # Drift
-    def f(self, t, y):
-        y = y.view(-1, *self.input_size)
-        return self.drift(t, y).flatten(start_dim=1) + self.score(t, y).flatten(start_dim=1)
-
-    # Diffusion
-    def g(self, t, y):
-        return torch.ones_like(y) * self.sigma