files and tweaks for SBI-sim notebook

2025-01-17 17:54:09 +08:00
1 parent 4c0fdd8dc0
commit 54d0dfc203
@@ -39,6 +39,7 @@ parts:
  - file: probmodels-intro.md
  - file: probmodels-ddpm-fm.ipynb
  - file: probmodels-phys.md
  - file: probmodels-sbisim.ipynb
  - file: probmodels-time.ipynb
  - file: probmodels-uncond.md
  - file: probmodels-graph.md
@@ -10,7 +10,9 @@
    "\n",
    "We'll be using a learning task where we can reliably generate arbitrary amounts of ground truth data, to make sure we can quantify how well the target distribution was learned. Specifically, we'll focus on Reynolds-averaged Navier-Stokes simulations around airfoils, which have the interesting characteristic that typical solvers (such as OpenFoam) transition from steady solutions to oscillating ones for larger Reynolds numbers. This transition is exactly what we'll give as a task to diffusion models below. (Details can be found in our [diffuion-based flow prediction repository](https://github.com/tum-pbs/Diffusion-based-Flow-Prediction/).) Also, to make the notebook self-contained, we'll revisit the most important concepts from the previous section.\n",
    "\n",
-    "> Note, if you're directly continuing reading from the previous chapter: an important difference to {doc}`probmodels-intro` is that, for simplicity, we'll apply denoising and flow-matching to a _forward_ problem here! We won't be aiming to recover $x$ for an observation $y$, but rather assume we have initial conditions $x$ from which we want to compute a solution $y$. So don't be surprised by the switched $x$ and $y$ below.\n",
+    "```{note} \n",
    "If you're directly continuing reading from the previous chapter, note that there's an important difference: for simplicity, we'll apply denoising and flow-matching to a **forward** problem here. We won't be aiming to recover $x$ for an observation $y$, but rather assume we have initial conditions $x$ from which we want to compute a solution $y$. So don't be surprised by the switched $x$ and $y$ below.\n",
    "```\n",
    "\n",
    "## Intro\n",
    "\n",
@@ -1,18 +1,43 @@
 Incorporating Physical Constraints
 =======================
-Despite all powerful capabilities of diffusion- and flow-based networks for generative modeling that we discussed in the previous sections, there is no direct feedback loop between the network, the observation and the sample at training time. This means there is no direct mechanism to include physics-based constraints such as priors from PDEs. As a consequence, it's very difficult to produce highly accurate samples based on learning alone: For scientific applications, we often want to make sure the errors go down to any chosen threshold.
+Despite the powerful capabilities of diffusion- and flow-based networks for generative modeling that we discussed in the previous sections, there is no direct feedback loop between the network, the observation and the sample at training time. This means there is no direct mechanism to include **physics-based constraints** such as priors from PDEs. As a consequence, it's very difficult to produce highly accurate samples based on learning alone: For scientific applications, we often want to make sure the errors go down to any chosen threshold.
 In this chapter, we will outline strategies to remedy this shortcoming, and building on the content of previous chapters, the central goal of both methods is to get **differentiable simulations** back into the training and inference loop. The previous chapters have shown that they're very capable tools, so the main question is how to best employ them in the context of diffusion modeling.
 ```{note} 
 Below we'll focus on the inverse problem setting from {doc}`probmodels-intro`. I.e., we have a system $y=f(x)$ (with numerical simulator $y=\mathcal P(x)$) and given an observation $y$, we'd like to obtain the posterior distribution for the distributional solution $x \sim p(x|y)$ of the inverse problem. 
 ```
 ## Guiding Diffusion Models
 Having access to a physical model with a differentiable simulation $\mathcal{P}(x)=y$ means we can obtain gradients $\nabla_x$ through the simulation. As before, we aim for solving _inverse_ problems where, given an output $y$ we'd like to sample from the conditional posterior distribution $p(x|y)$ to obtain samples $x$ that explain $y$. The previous chapter demonstrated learning such distributions with diffusion models, and given a physics prior $\mathcal{P}$, there's a first fundamental choice: should be use the gradient at _training time_, i.e., trying to improve the learned distribution $p_\theta$, or at _inference time_, to improve sampling $x \sim p_\theta(x|y)$.
 **Training with physics priors:** The hope of incorporating physics-based signals in the form of gradients at training time would be to improve the state of $p_\theta$ after training. While there's a certain hope this could, e.g., compensate for sparse training data, there is little hope for substantially improving the accuracy of the learned distribution. The training process for diffusion and flow matching models typically yields very capable neural networks, that are excellent at producing approximate samples from the posterior. They're typically limited in terms of their accuracy by model and training data size, but it's difficult to fundamentally improve the capabilities of a model at this stage. Rather, in this context it is more interesting to obtain higher accuracies at inference time.
 **Inference with physics priors:** For scientific applications, classic simulations typically yield control knobs that allow for choosing a level of accuracy. E.g., iterative solvers for linear systems provide iteration counts and residual thresholds, and if a solution is not accurate enough, a user can simply reduce the residual threshold to obtain a more accurate output. In contrast, neural networks typically come without such controls, and even the iteration count of denoising or velocity integration (for flow matching) are bounded in terms of final accuracy. More steps typically reduce noise, and correspondingly the error, but will plateau at a level of accuracy given by the capabilities of the trained model. This is exactly where the gradients of physics solver show promise: they provide an external process that can guide and improve the output of a diffusion model. As we'll show below, this makes it possible to push the levels of accuracy beyond those of pure learning, and can yield inverse problem solvers that really outperform traditional solvers.
 Recall that for denoising, we train a noise estimator $\epsilon_\theta$, and at inference time iterate denoising steps of the form
 $x_{\text{new}} =  x - \hat \alpha_t \epsilon_\theta(x, t) + \hat \sigma_t \mathcal N(0,I)$ , where $\hat \alpha,\hat \sigma$ denote the merged scaling factors for both terms.
 The most straight-forward approach for including gradients is to additionally include a step in the direction of the gradient $\nabla_x || \mathcal P(x) - y||_2$. For simplicity, we take an $L^2$ distance towards the observation $y$ here. This was shown to direct sampling even when the posterior is not conditional, i.e., if we only have access to $x \sim p_\theta(x)$, and is known as _diffusion posterior sampling_ {cite}`chung2023dps`.
 While this approach manages to includes $\mathcal P$, there are two challenges: $x$ is typically noisy, and the gradient step can distort the distributional sampling of the denoising process. The first point is handled quite easily with an _extrapolation step_ (more details below), while the second one is more difficult to address: the gradient descent steps via $\nabla_x \mathcal P$ are akin to a classic optimization for the inverse problem and could strongly distort the outputs of the diffusion model. E.g., in the worst case they could pull the different points of the posterior distribution towards a single case favored by the simulator $\mathcal P$. Hence, the following paragraphs will outline a strategy that merges simulator and learning, while preserving the distribution of the posterior.
 We'll focus on flow matching as a state-of-the-art approach next, and afterwards discuss variant that treats the diffusion steps themselves as a physical process.
 In this chapter, we will outline two strategies to remedy this shortcoming, and building on the content of previous chapters, the central goal of both methods is to get **differentiable simulations** back into the training and inference loop. The previous chapters have shown that they're very capable tools, so the main question is how to bext employ them in the context of diffusion modeling.
 ![Divider](resources/divider5.jpg)
 ## Physics-Guided Flow Matching
-First, we'll outline a strategy to reintroduce control signals using simulators into the flow matching algorithm {cite}`holzschuh2024fm`. We transform an existing pretrained flow-based network, as outlined in {doc}`probmodels-intro`, with a flexible 
+To reintroduce control signals using simulators into the flow matching algorithm we'll follow {cite}`holzschuh2024fm`. The goal is to transform an existing pretrained flow-based network, as outlined in {doc}`probmodels-intro`, with a flexible control signal by aggregating the learned flow and control signals into a _controlled flow_. This is the task of a second neural network, the _control network_, in  order to make sure that the posterior distribution is not negatively affected by the signals from the simulator. This second network is small compared to the pretrained flow network, and freezing the weights of the pretrained network works very well; thus, the refinement for control needs only a fairly small amount of additional parameters and computing resources.
 control signal by aggregating the learned flow and control signals into a _controlled flow_. The aggregation network is small compared to the pretrained flow network, and we find that freezing the weights of the pretrained network works very well; thus, refining needs only a minimal amount of additional parameters and compute.
 ```{figure} resources/probmodels-phys-overview.jpg
@@ -20,7 +45,7 @@ control signal by aggregating the learned flow and control signals into a _contr
 height: 240px
 name: probmodels-phys-overview
 ---
-An overview of the control framework. We will consider a pretrained flow network $v_\phi$ and use the predicted flow for the trajectory point $x_t$ at time $t$ to estimate $\hat{x}_1$.
+An overview of the control framework. We will consider a pretrained flow network $v_\theta$ and use the predicted flow for the trajectory point $x_t$ at time $t$ to estimate $\hat{x}_1$.
 On the right, we show a gradient-based control signal with a differentiable simulator and cost function $C$ for improving $\hat{x}_1$.
 An additional network learns to combine the predicted flow with feedback via the control signal to give a new controlled flow.
 By combining learning-based updates with suitable controls, we avoid local optima and obtain high-accuracy samples with low inference times.
@@ -28,34 +53,30 @@ By combining learning-based updates with suitable controls, we avoid local optim
 The control signals can be based on gradients and a cost function, if the simulator is differentiable, but they can also be learned directly from the simulator output.
 Below, we'll show that performance gains due to simulator feedback are substantial and cannot be achieved by training on larger datasets alone.
-Specifically, we'll show that flow matching with simulator feedback is competitive with MCMC baselines for a problem from gravitational in terms of accuracy, and it beats them significantly regarding inference time. This indicates that it poses a very attractive tool for practical applications.
+Specifically, we'll show that flow matching with simulator feedback is competitive with MCMC baselines for a problem from gravitational lensing in terms of accuracy, and it beats them significantly regarding inference time. This indicates that it provides a very attractive tool for practical applications.
 Note that in the following we'll focus on the inverse problem setting from {doc}`probmodels-intro`. I.e., we have a system $y=f(x)$ and given an observation $y$, we'd like to obtain the posterior distribution for the distributional solution $x \sim p(x|y)$ of the inverse problem. 
-**1-step prediction** 
+**Controlled flow $v_\theta^C$** First, it's a good idea to pretrain a regular, conditional flow network $v_\theta(x,y,t)$ without any control signals to make sure that we can realize the best achievable performance possible based on learning alone. 
-The flow matching networks $v_\phi(t,x)$ from {doc}`probmodels-intro` gradually transform samples from $p_0$ to $p_1$ during inference via integrating a simple ODE step by step. There is no direct feedback loop between the current point on the trajectory $x_t$, the observation $y$, and a physical model that we could bring into the picture. An important first issue is that the current trajectory point $x_t$ is often not be close to a good estimate of a posterior sample $x_1$. 
+Then, in a second training phase, a control network $v_\theta^C(v, c,t)$ is introduced. It receives the pretrained flow $v$ and control signal $c$ as input. Based on these additional inputs, it can used, e.g., the gradient of a PDE to produce an improved flow matching velocity. At inference time, we integrate
 $dx/dt = v^C_\theta(v,c,t)$ just like before, only now this means evaluating $v_\theta(x,y,t)$ and then $c$ beforehand. (We'll focus on the details of $c$ in a moment.)
 First, the control network is much smaller in size than the regular flow network, making up ca. $10\%$ of the weights $\theta$. The network weights of $v_\theta$ can be frozen, to train with the conditional flow matching loss {eq}`conditional-flow-matching` for a small number of additional steps. This reduces training time and compute since we do not need to backpropagate gradients through $v_\theta(x, y,t)$. Freezing the weights of $v_\theta$ typically does not negatively affects the performance, although a joint end-to-end training could provide some additional improvements. 
 **1-step prediction** The conditional flow matching networks $v_\theta(x,y,t)$ from {doc}`probmodels-intro` gradually transform samples from $p_0$ to $p_1$ during inference via integrating the simple ODE $dx_t/dt = v_\theta(x_t,y,t)$ step by step. There is no direct feedback loop between the current point on the trajectory $x_t$, the observation $y$, and a physical model that we could bring into the picture. An important first issue is that the current trajectory point $x_t$ is often not be close to a good estimate of a posterior sample $x_1$. 
 This is especially severe at the beginning of inference, where $x_0$ is drawn from the source distribution (typically a Gaussian), and hence $x_t$ will be very noisy. Most simulators really don't like very noisy inputs, and trying to compute gradients on top of it is clearly a very bad idea.
 This issue is alleviated by extrapolating $x_t$ forward in time to obtain an estimated $\hat{x}_1$ 
 $$
 \begin{align} 
-    \hat{x}_1 = x_t + (1-t) v_\phi(t, x_t, y).
+    \hat{x}_1 = x_t + (1-t) v_\theta(x_t, y, t).
 \end{align}
 $$ (eq:1_step_prediction)
 and then performing subsequent operations for control and guidance on $\hat{x}_1$ instead of the current, potentially noisy $x_1$.
 ```{note} Direct Gradient Descent.
 not great, similar to DPS?
 uncond diffusion + gradients
 ```
 Note that this 1-step prediction is also conceptually related to diffusion sampling using [_likelihood-guidance_](http://DBLP:conf/nips/WuTNBC23). For inference in diffusion models, where sampling is based on the conditional score $\nabla_{x_t} \log p(x_t|y)$ and can be decomposed into 
 $$
@@ -69,68 +90,58 @@ $p(y|x_t) \approx p_{y|x_0}(y|\hat{x}(x_t))$,
 where the denoising estimate $\hat{x}(x_t) = \mathbb{E}_q[x_0|x_t]$ is usually obtained via Tweedie's formula $(\mathbb{E}_q[x_0|x_t] - x_t) / t\sigma^2$. In practice, the estimate $\hat{x}(x_t)$ is very poor when $x_t$ is still noisy, making inference difficult in the early stages. In contrast, flows based on linear conditional transportation paths have empirically been shown to have trajectories with less curvature compared to, for example, denoising-based networks. This property of flow matching enables inference in fewer steps and providing better estimates for $\hat{x}_1$.
 **Controlled flow $v_\phi^C$** 
 First, it's a good idea to pretrain a regular flow network $v_\phi(x,y,t)$ without any control signals to make sure that we can realize the best achievable performance possible based on learning alone. 
 Then, in a second training phase, a control network $v_\phi^C(v, c,t)$ is introduced. It receives the pretrained flow $v$ and control signal $c$ as input. Based on these additional inputs, it can used, e.g., the gradient of a PDE to produce an improved flow matching velocity. At inference time, we integrate
 $dx/dt = v^C_\phi(v,c,t)$ just like before, only now this means evaluating $v_\phi(x,y,t)$ and then $c$ beforehand. (We'll focus on the details of $c$ in a moment.)
 First, the control network is much smaller in size than the regular flow network, making up ca. $10\%$ of the weights $\phi$. The network weights of $v_\phi$ can be frozen, to train with the conditional flow matching loss {eq}`conditional-flow-matching` for a small number of additional steps. This reduces training time and compute since we do not need to backpropagate gradients through $v_\phi(x, y,t)$. We did not observe that freezing the weights of $v_\phi$ affects the performance negatively. We include algorithms for training in appendix \ref{app:algorithms}.
-## Physics-based Controls
+
 ### Physics-based Controls
 Now we focus on the content of the control signal $c$ that was already used above. We extend the idea of self-conditioning via physics-based control signals to include an additional feedback loop between the network output and an underlying physics-based prior. We'll distinguish between two types of controls in the following: a gradient-based control from a differentiable simulator, and one from a learned estimator network. 
-```{figure} resources/probmodels-learning_based_control_signal.jpg
+```{figure} resources/probphys02-control.jpg
 ---
 height: 240px
-name: learning_based_control_signal
+name: probphys02-control
 ---
-TODO, learning_based_control_signal figure
+Types of control signals. (a) From a differentiable simulator, and (b) from a learned encoder.
 ```
-**Gradient-based control signal**
+**Gradient-based control signal** In the first case, we make use of a differentiable simulator $\mathcal{P}$ to construct a cost function $C$. Naturally, $C$ will likewise be differentiable such that we can compute a gradient for a predicted solution. Also, we will rely on the stochasticity of diffusion/flow matching, and as such the simulator can be deterministic.
 In the first case, we make use of a differentiable simulator $S$ to construct a cost function $C$. Naturally, $C$ will likewise be differentiable such that we can compute a gradient for a predicted solution. Also, we will rely on the stochasticity of diffusion/flow matching, and as such the simulator can be deterministic.
 Given an observation $y$ and the estimated 1-step prediction $\hat{x}_1$, the control signal computes to how well $\hat{x}_1$ explains $y$ via the cost function $C$. Good choices for the cost are, e.g., an $L^2$ loss or a likelihood $p(y|\hat{x}_1)$. We define the control signal $c$ to consist of two components: the cost itself, and the gradient w.r.t. the cost function:  
 $$
    \begin{align}
-    c(\hat{x}_1, y) := [C(S(\hat{x}_1), y); \nabla_{\hat{x}_1} C(S(\hat{x}_1), y)].
+    c(\hat{x}_1, y) := [C(\mathcal{P}(\hat{x}_1), y); \nabla_{\hat{x}_1} C(\mathcal{P}(\hat{x}_1), y)].
    \end{align}
 $$
 As this information is passed to a network, the network can freely make use of the current distance to the target (the value of $C$) and the direction towards lowering it in the form of $\nabla_{\hat{x}_1} C$.
-**Learning-based control signal**
+**Learning-based control signal** When the simulator is non-differentiable, the second variant of using a learned estimator comes in handy. 
-
+To combine the simulator output with the observation $y$, a learnable encoder network _Enc_ with parameters $\theta_E$ can be introduced to judge the similarity of the simulation and the observation. The output of the encoder is small and of size $O(\mathrm{dim}(x))$.
 When the simulator is non-differentiable, the second variant of using a learned estimaor comes in handy. 
 To combine the simulator output with the observation $y$, a learnable encoder network _Enc_ with parameters $\phi_E$ can be introduced to judge the similarity of the simulation and the observation. The output of the encoder is small and of size $O(\mathrm{dim}(x))$.
 The control signal is then defined as 
 $$
 \begin{align}
-    c(\hat{x}_1, y) := Enc(S(\hat{x}_1), y).
+    c(\hat{x}_1, y) := Enc(\mathcal{P}(\hat{x}_1), y).
 \end{align}
 $$
-The gradient backpropagation is stopped at the output of the simulator $S$, as shown in {ref}`learning_based_control_signal`. 
+The gradient backpropagation is stopped at the output of the simulator $\mathcal{P}$, as shown in {numref}`figure {number} <probphys02-control>`. 
 Before showing some examples of the capabilities of these two types of control, we'll discuss some of their properties.
-## Additional Considerations 
+### Additional Considerations 
-**Stochastic simulators**
+**Stochastic simulators** Many Bayesian inference problems have a stochastic simulator. For simplicity, we assume that all stochasticity within such a simulator can be controlled via a variable $z \sim \mathcal{N}(0, I)$, which is an additional input. Motivated by the equivalence of exchanging expectation and gradient  
 Many Bayesian inference problems have a stochastic simulator. For simplicity, we assume that all stochasticity within such a simulator can be controlled via a variable $z \sim \mathcal{N}(0, I)$, which is an additional input. Motivated by the equivalence of exchanging expectation and gradient  
 $$
 \begin{align}
-    \nabla_{\hat{x}_1} \mathbb{E}_{z\sim \mathcal{N}(0,1)} [ C(S_z(\hat{x}_1), y)] = \mathbb{E}_{z\sim \mathcal{N}(0,1)} [ \nabla_{\hat{x}_1} C(S_z(\hat{x}_1), y)],
+    \nabla_{\hat{x}_1} \mathbb{E}_{z\sim \mathcal{N}(0,1)} [ C(\mathcal P_z(\hat{x}_1), y)] = \mathbb{E}_{z\sim \mathcal{N}(0,1)} [ \nabla_{\hat{x}_1} C(\mathcal P_z(\hat{x}_1), y)],
 \end{align}
 $$
@@ -140,27 +151,136 @@ when calling the simulator, we draw a random realization of $z$. During training
 If the estimate $\hat{x}_1$ is bad and the corresponding cost $C(\hat{x}_1, y)$ is high, gradients and control signals can become unreliable. It turns out that the estimates $\hat{x}_1$ become more reliable for later times in the flow matching process. 
-In practice, $t \geq 0.8$ is a good threshold. Therefore, we only train the control network $v_\phi^C$ in this range, which allows for focusing on control signals containing more useful information to, e.g. fine tune the solutions with the accurate gradients of a differentiable simulator. For $t < 0.8$, we directly output the pretrained flow $v_\phi(t, x, y)$.
+In practice, $t \geq 0.8$ is a good threshold. Therefore, we only train the control network $v_\theta^C$ in this range, which allows for focusing on control signals containing more useful information to, e.g. fine tune the solutions with the accurate gradients of a differentiable simulator. For $t < 0.8$, we directly output the pretrained flow $v_\theta(t, x, y)$.
 **Theoretical correctness**
-In the formulation above, the approximation $\hat{x}_1$ only influences the control signal, which is an input to the controlled flow network $v_\phi^C$. In the case of a deterministic simulator, this makes the control signal a function of $x_t$. The controlled flow network is trained with the same loss as vanilla flow matching. This has the nice consequence that the theoretical properties are preserved.
+In the formulation above, the approximation $\hat{x}_1$ only influences the control signal, which is an input to the controlled flow network $v_\theta^C$. In the case of a deterministic simulator, this makes the control signal a function of $x_t$. The controlled flow network is trained with the same loss as vanilla flow matching. This has the nice consequence that the theoretical properties are preserved.
 This is in contrast to e.g. "likelihood-based guidance", which uses an approximation for $\nabla_{x_t} \log p(y|x_t)$ as a guidance term during inference, which is not covered by the original flow matching theory. 
-## An Example from Astrophysics
+### An Example from Astrophysics
-To demonstrate how these guidance from a physics solver affect the accuracy of samples and the posterior, we show an example from strong gravitational lensing:  an inverse problem in astrophysics that is challenging and requires precise posteriors for accurate modeling of observations. In galaxy-scale strong lenses, light from a source galaxy is deflected by the gravitational potential of a galaxy between the source and observer, causing multiple images of the source to be seen. Traditional computational approaches require several minutes to many hours or days to model a single lens system. Therefore, there is an urgent need to reduce the compute and inference with learning-based methods. In this experiment, we demonstrate that using flow matching and the control signals with feedback from a simulator, we obtain posterior distributions for lens modeling that are competitive with the posteriors obtained by MCMC-based methods but with much faster inference times. 
+To demonstrate how these guidance from a physics solver affect the accuracy of samples and the posterior, we show an example from strong gravitational lensing:  an inverse problem in astrophysics that is challenging and requires precise posteriors for accurate modeling of observations. In galaxy-scale strong lenses, light from a source galaxy is deflected by the gravitational potential of a galaxy between the source and observer, causing multiple images of the source to be seen. Traditional computational approaches require several minutes to many hours or days to model a single lens system. Therefore, there is an urgent need to reduce the compute and inference with learning-based methods. In this experiment, it's shown that using flow matching and the control signals with feedback from a simulator gives posterior distributions for lens modeling that are competitive with the posteriors obtained by MCMC-based methods. At the same time, they are much faster at inference. 
-...
+
 ```{figure} resources/probmodels-astro.jpg
 ---
 height: 240px
 name: probmodels-astro
 ---
 Results from flow matching for reconstructing gravitational lenses. Left: flow matching with a differentiable simulator (bottom) clearly outperforms pure flow matching (top). Right: comparisons against classic baselines. The FM+simulator variant is more accurate while being faster.
 ```
 The image aboves shows an example reconstruction and the residual errors. While flow matching and the physics-based variant are both very accurate (it's hard to visually make out differences), the FM version is just on par with classic inverse solvers. The version with the simulator, however, provides a substantial boost in terms of accuracy that is very difficult to achieve even for classic solvers. The quantitative results are shown in the table on the right: the best classic baseline is AIES with an average $\chi_2$ statistic of 1.74, while FM with simulator yields 1.48. Provided that the best possible result due to noisy observations is 1.17 for this scenario, the FM+simulation version is really highly accurate. 
 At the same time, the performance numbers for _modeling time_ in the right column show that the FM variant clearly outperforms the classic solvers. While the simulator increases inference time compared to only the neural network (10s to 19s), the classic baselines require more than $50\times$ longer reconstruction times. Interestingly, this example also highlights the problems of "simpler" physics combinations in the form of DPS. The DPS version does not manage to keep up with the classic solvers in terms of accuracy. To conclude, the _FM+simulator_ variant is not only substantially more accurate, but also ca. $35\times$ faster than the best classic solver above (AIES). (Source code for this approach will be available soon [in this repository](https://github.com/tum-pbs/sbi-sim).)
 ---
 A summary of the physics-based flow matching is given by the following bullet points:
 ✅ Pro:
 * Improved accuracy over purely learned diffusion models
 * Gives control over residual accuracy
 * Reduced runtime compared to traditional inverse solvers
 ❌ Con:
 * Requires differentiable physical process
 * Increased computational resources
 ![Divider](resources/divider6.jpg)
 ## Score Matching with Differentiable Physics
-alternative: equate diffusion time and physics
+So far we have treated the _diffusion time_ of denoising and flow matching as a process that is purely virtual and orthogonal to the time of the physical process to be represented by the forward and inverse problems. This is the most generic viewpoint, and works nicely, as demonstrated above. However, it's interesting to think about the alternative: merging the two processes, i.e., treating the diffusion process as an inherent component of the physics system.
-sbi-sim ...?
+```{figure} resources/probmodels-smdp-1trainB.jpg
 ---
 height: 240px
 name: probmodels-smdp-trainB
 ---
 The physics process (heat diffusion as an example, left) perturbs and "destroys" the initial state. At inference time (right, Buoyancy flow as an example), the solver is used to compute inverse steps and produce solutions by combining steps along the score and the gradient of the solver.
 ```
 The following sections will explain such a combined approach, following the paper "Solving Inverse Physics Problems with Score Matching" {cite}`holzschuh2023smdp`, which which [code is available in this repository](https://github.com/tum-pbs/SMDP).
 This approach solves inverse physics problems by leveraging the ideas of score matching. The system’s current state is moved backward in time step by step by combining an approximate inverse physics simulator and a learned correction function. A central insight of this work is that training the learned correction with a single-step loss is equivalent to a score matching objective, while recursively predicting longer parts of the trajectory during training relates to maximum likelihood training of a corresponding probability flow. The resulting inverse solver exhibits good accuracy and temporal stability. In line with diffusion modeling and in contrast to classic learned solvers, it allows for sampling the posterior of the solutions. The method will be called _SMDP_ (for _Score Matching with Differentiable Physics_) in the following.
 ### Training and Inference with SMDP
 For training, SMDP fits a neural ODE, the probability flow, to the set of perturbed training trajectories. The probability flow is comprised of an approximate reverse physics simulator $\tilde{\mathcal{P}}^{-1}$ as well as a correction function $s_\theta$. For inference, we simulate the system backward in time from $\mathbf{x}_T$ to $\mathbf{x}_0$ by combining $\tilde{\mathcal{P}}^{-1}$, the trained $s_\theta$ and Gaussian noise in each step. 
 For optimizing $s_\theta$, our approach moves a sliding window of size $S$ along the training trajectories and reconstructs the current window. Gradients for $\theta$ are accumulated and backpropagated through all prediction steps. This process is illustrated in the following figure:
 ```{figure} resources/probmodels-smdp-1train.jpg
 ---
 height: 240px
 name: probmodels-smdp-train
 ---
 Overview of the score matching training process while incorporating a physics solver $\mathcal P$ and it's approximate inverse solver $\matcal{P}^{-1}.
 ```
 A differentiable solver or a learned surrogate model is employed for $\tilde{\mathcal{P}}^{-1}$. 
 The neural network $s_\theta(\mathbf{x}, t)$ parameterized by $\theta$ is trained such that
 $$
 \mathbf{x}_{m} \approx \mathbf{x}_{m+1} + \Delta t \left[ \tilde{\mathcal{P}}^{-1}(\mathbf{x}_{m+1}) + s_\theta(\mathbf{x}_{m+1}, t_{m+1}) \right].
 $$
 In this equation, the term $s_\theta(\mathbf{x}_{m+1}, t_{m+1})$ corrects approximation errors and resolves uncertainties from the stochastic forcing $F_{t_m}(z_m)$. Potentially, this process can be unrolled over multiple steps at training time to improve accuracy and stability. At inference, time the stochastic differential equation 
 $$
 d\mathbf{x} = \left[ -\tilde{\mathcal{P}}^{-1}(\mathbf{x}) + C \, s_\theta(\mathbf{x},t) \right] dt + g(t) dW
 $$
 is integrated via the Euler-Maruyama method to obtain a solution for the inverse problem.
 Setting $C=1$ and excluding the noise gives the probability flow ODE: a unique, deterministic solution. This deterministic variant is not probablistic anymore, but has other interesting properties. 
 ```{figure} resources/probmodels-smdp-2infer.jpg
 ---
 height: 148px
 name: probmodels-smdp-infer
 ---
 An overview of SMDP at inference time.
 ```
 ### SMDP in Action
 This section shows experiments for the stochastic heat equation: $\frac{\partial u}{\partial t} = \alpha \Delta u$, which plays a fundamental role in many physical systems. It slightly perturbs the heat diffusion process and includes an additional term $g(t)\ \xi$, where $\xi$ is space-time white noise. For the experiments, we fix the diffusivity constant to $\alpha = 1$ and sample initial conditions at $t=0$ from Gaussian random fields with $n=4$ at resolution $32 \times 32$. We simulate the heat diffusion with noise from $t=0$ until $t=0.2$ using the Euler-Maruyama method and a spectral solver $\mathcal{P}_h$ with a fixed step size and $g \equiv 0.1$. Given a simulation end state $\mathbf{x}_T$, we want to recover a possible initial state $\mathbf{x}_0$. 
 In this experiment, the forward solver cannot be used to infer $\mathbf{x}_0$ directly since high frequencies due to noise are amplified, leading to physically implausible solutions. Instead, the reverse physics step $\tilde{P}^{-1}$ is implemented by using the forward step of the solver $\mathcal{P}_h(\mathbf{x})$, i.e. $\tilde{\mathcal{P}}^{-1}(\mathbf{x}) \approx - \mathcal{P}_h (\mathbf{x})$.
 A small ResNet-like architecture is used based on an encoder and decoder part as representation for the score function $s_\theta(\mathbf{x}, t)$. The spectral solver is implemented via differentiable programming in _JAX_. As baseline methods, a supervised training of the same architecture as $s_\theta(\mathbf{x}, t)$, a Bayesian neural network (BNN), as well as a FNO network are considered. An $L_2$ loss is used for all these methods, i.e., the training data consists of pairs of initial state $\mathbf{x}_0$ and end state $\mathbf{x}_T$. Additionally, a variant of the SMDP method is included for which the reverse physics step $\tilde{\mathcal{P}}^{-1}$ is reomved, such that the inversion of the dynamics has to be learned entirely by $s_\theta$, denoted by ''$s_\theta$~only''.
 ```{figure} resources/probmodels-smdp-3heat.jpg
 ---
 name: probmodels-smdp-heat
 ---
 While the ODE trajectories provide smooth solutions with the lowest reconstruction MSE, the SDE solutions synthesize high-frequency content, significantly improving spectral error.  
 The ``$s_\theta$ only'' version without the reverse physics step exhibits a significantly larger spectral error. Metrics (right) are averaged over three runs.
 ```
 SMDP and the baselines are evaluated by considering the _reconstruction MSE_ on a test set of $500$ initial conditions and end states. For the reconstruction MSE, the prediction of the network is simulated forward in time with the solver $\mathcal{P}_h$ to obtain a corresponding end state, which is compared to the ground truth via the $L_2$ distance. This metric has the disadvantage that it does not measure how well the prediction matches the training data manifold. I.e., for this case, whether the prediction resembles the properties of the initial Gaussian random field. For that reason, the power spectral density of the states is shown as a _spectral loss_. An evaluation and visualization of the reconstructions are given in figure \ref{fig:stochastic_heat_eq_overview}, which shows that the ODE inference performs best regarding the reconstruction MSE. However, its solutions are smooth and do not contain the necessary small-scale structures. This is reflected in a high spectral error. The SDE variant, on the other hand, performs very well in terms of spectral error and yields visually convincing solutions with only a slight increase in the reconstruction MSE. 
 This highlights the role of noise as a source of entropy in the inference process for diffusion models, such as the SDE in SMDP, which is essential for synthesizing small-scale structures. Note that there is a natural tradeoff between both metrics, and the ODE and SDE inference perform best for each of the cases while using an identical set of weights. This heat diffusion example highlights the advantages and properties of treating the physical process as part of the diffusion process. This, of course, extends to other physics. E.g., [the SMDP repository](https://github.com/tum-pbs/SMDP) additionally shows a case with an inverse Navier-Stokes solve.
 ## Summary of Physics-based Diffusion Models
 Overall, the sections above have explained two methods to incorporate physics-based constraints and models in the form of PDEs into diffusion modeling. Interestingly, the inclusion is largely in line with {doc}`diffphys`, i.e. gradients of the physics solver are a central quantity, and concepts like unrolling play an important role. On the other hand, the probabilistic modeling introduces additional complexity on the training and inference sides. It provides powerful tools and access to distribiutions of solutions (we haven't even touched follow up applications such as uncertainty quantification above), but this comes at a cost. 
 As a rule of thumb 👍, diffusion modeling should only be used if the solution is a distribution that is _not_ well represented by the mean of the solutions. If the mean is accetable, "regular" neural networks offer substantial advantages in terms of reduced complexity for training and inference.
 However, if the solutions are a distribution 🌦️, diffusion models are powerful tools to work with complex and varied solutions. Given its capabilties, deep learning with diffusion models arguably introduces surprisingly _little_ additional complexity. E.g., training flow matching models is suprisingly robust, can be build on top of deterministic training, and introduces only a mild computational overhead.
 To show how the combination of physics solvers and diffusion models turns out in terms of an implementation, the next section shows source code for an SMDP use case.
@@ -7,12 +7,38 @@
        "id": "4139ba7c-6234-47b8-8a37-35cd01efd51d"
      },
      "source": [
-        "# Inverse Problem Example with Differentiable Simulations\n",
+        "Inverse Problem Example with Differentiable Simulations\n",
        "=======================\n",
        "\n",
-        "The purpose of this notebook is to provide a simple implementation of our proposed SMDP method, which combines physics and score matching to solve inverse problems. </br> We consider the simple toy problem, which is shown in Figure 1 in the main paper and discussed in more detail in Appendix B. </br>\n",
+        "This notebook will illustrate some of the concepts introduced in {doc}`probmodels-intro`, such as the training of score functions via log likelihoods, and what they look like in a clear and reduced problem. At the same time, the setup provides integration of a simple _differentiable simulator_ to illustrate the concept of physics-based diffusion modeling with the SMDP method from {doc}`probmodels-phys` [(full paper)](https://arxiv.org/abs/2301.10250). This approach combines physics and score matching along a merged time dimension to solve inverse problems. \n",
        "\n",
-        "Our implementation comprises the data generation, network definition, training using the sliding window method, </br>\n",
+        "## Toy Problem setup\n",
-        "visualization of the learned score and inference using the probability flow ODE and the reverse-time SDE. </br>\n",
+        "\n",
        "We'll consider a toy problem with quadratically decaying trajectories. The trajectories start at $1$ or $-1$ and approach $0$ as $t$ increases. \n",
        "The corresponding SDE is given by \n",
        "\n",
        "$$\n",
        "\\begin{align} \n",
        "    dx = - \\left[ \\lambda_1 \\cdot \\mathrm{sign}(x) x^2 \\right] dt + \\lambda_2 dw,\n",
        "\\end{align}\n",
        "$$ (eq:app_sde_1d)\n",
        "\n",
        "with $\\lambda_1 = 7$ and $\\lambda_2 = 0.03$. The corresponding reverse-time SDE is \n",
        "\n",
        "$$\n",
        "\\begin{align}\n",
        "    dx = - \\left[ \\lambda_1 \\cdot \\mathrm{sign}(x) x^2 - \\lambda_2^2 \\cdot \\nabla_x \\log p_t(x) \\right] dt + \\lambda_2 dw.\n",
        "\\end{align} \n",
        "$$\n",
        "\n",
        "Throughout the experiments, $p_0$ is a categorical distribution, where we draw either $1$ or $-1$ with the same probability. \n",
        "\n",
        "\n",
        "\n",
        "## Implementation Overview\n",
        "\n",
        "The implementation below comprises the data generation, network definition, training using the sliding window method, \n",
        "visualization of the learned score and inference using the probability flow ODE and the reverse-time SDE. \n",
        "The core algorithm for training and inference can be modified easily for different variants that are discussed in the paper."
      ]
    },
@@ -23,13 +49,13 @@
        "id": "3d365122-f676-425f-8445-39ceadfcbcd5"
      },
      "source": [
-        "### Required packages\n",
+        "\n",
-        "Our training implementation uses **JAX** with **haiku** and **optax** as neural network libraries. For data generation and inference, we use **diffrax**==0.1.2"
+        "This implementation uses quite a few packages, which we'll import first. We'll make use of `JAX` with `haiku` and `optax` as neural network libraries. For data generation and inference, we use `diffrax`:"
      ]
    },
    {
      "cell_type": "code",
-      "execution_count": 1,
+      "execution_count": null,
      "id": "0_JufNgq_mMQ",
      "metadata": {
        "colab": {
@@ -38,83 +64,14 @@
        "id": "0_JufNgq_mMQ",
        "outputId": "88b5282d-d062-4d60-d26a-0bc533da4774"
      },
-      "outputs": [
+      "outputs": [],
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Collecting diffrax\n",
            "  Downloading diffrax-0.6.2-py3-none-any.whl.metadata (17 kB)\n",
            "Requirement already satisfied: jax in /usr/local/lib/python3.10/dist-packages (0.4.33)\n",
            "Requirement already satisfied: jaxlib in /usr/local/lib/python3.10/dist-packages (0.4.33)\n",
            "Requirement already satisfied: scipy in /usr/local/lib/python3.10/dist-packages (1.13.1)\n",
            "Requirement already satisfied: optax in /usr/local/lib/python3.10/dist-packages (0.2.4)\n",
            "Collecting dm-haiku\n",
            "  Downloading dm_haiku-0.0.13-py3-none-any.whl.metadata (19 kB)\n",
            "Collecting equinox>=0.11.10 (from diffrax)\n",
            "  Downloading equinox-0.11.11-py3-none-any.whl.metadata (18 kB)\n",
            "Collecting jax\n",
            "  Downloading jax-0.4.38-py3-none-any.whl.metadata (22 kB)\n",
            "Collecting jaxtyping>=0.2.24 (from diffrax)\n",
            "  Downloading jaxtyping-0.2.36-py3-none-any.whl.metadata (6.5 kB)\n",
            "Collecting lineax>=0.0.5 (from diffrax)\n",
            "  Downloading lineax-0.0.7-py3-none-any.whl.metadata (17 kB)\n",
            "Collecting optimistix>=0.0.7 (from diffrax)\n",
            "  Downloading optimistix-0.0.10-py3-none-any.whl.metadata (17 kB)\n",
            "Collecting typeguard==2.13.3 (from diffrax)\n",
            "  Downloading typeguard-2.13.3-py3-none-any.whl.metadata (3.6 kB)\n",
            "Requirement already satisfied: typing-extensions>=4.5.0 in /usr/local/lib/python3.10/dist-packages (from diffrax) (4.12.2)\n",
            "Collecting jaxlib\n",
            "  Downloading jaxlib-0.4.38-cp310-cp310-manylinux2014_x86_64.whl.metadata (1.0 kB)\n",
            "Requirement already satisfied: ml_dtypes>=0.4.0 in /usr/local/lib/python3.10/dist-packages (from jax) (0.4.1)\n",
            "Requirement already satisfied: numpy>=1.24 in /usr/local/lib/python3.10/dist-packages (from jax) (1.26.4)\n",
            "Requirement already satisfied: opt_einsum in /usr/local/lib/python3.10/dist-packages (from jax) (3.4.0)\n",
            "Requirement already satisfied: absl-py>=0.7.1 in /usr/local/lib/python3.10/dist-packages (from optax) (1.4.0)\n",
            "Requirement already satisfied: chex>=0.1.87 in /usr/local/lib/python3.10/dist-packages (from optax) (0.1.88)\n",
            "Requirement already satisfied: etils[epy] in /usr/local/lib/python3.10/dist-packages (from optax) (1.11.0)\n",
            "Collecting jmp>=0.0.2 (from dm-haiku)\n",
            "  Downloading jmp-0.0.4-py3-none-any.whl.metadata (8.9 kB)\n",
            "Requirement already satisfied: tabulate>=0.8.9 in /usr/local/lib/python3.10/dist-packages (from dm-haiku) (0.9.0)\n",
            "Requirement already satisfied: toolz>=0.9.0 in /usr/local/lib/python3.10/dist-packages (from chex>=0.1.87->optax) (0.12.1)\n",
            "Downloading diffrax-0.6.2-py3-none-any.whl (187 kB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m187.2/187.2 kB\u001b[0m \u001b[31m4.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hDownloading typeguard-2.13.3-py3-none-any.whl (17 kB)\n",
            "Downloading jax-0.4.38-py3-none-any.whl (2.2 MB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m2.2/2.2 MB\u001b[0m \u001b[31m37.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hDownloading jaxlib-0.4.38-cp310-cp310-manylinux2014_x86_64.whl (101.7 MB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m101.7/101.7 MB\u001b[0m \u001b[31m21.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hDownloading dm_haiku-0.0.13-py3-none-any.whl (373 kB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m373.9/373.9 kB\u001b[0m \u001b[31m16.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hDownloading equinox-0.11.11-py3-none-any.whl (179 kB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m179.2/179.2 kB\u001b[0m \u001b[31m10.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hDownloading jaxtyping-0.2.36-py3-none-any.whl (55 kB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m55.8/55.8 kB\u001b[0m \u001b[31m3.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hDownloading jmp-0.0.4-py3-none-any.whl (18 kB)\n",
            "Downloading lineax-0.0.7-py3-none-any.whl (67 kB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m67.3/67.3 kB\u001b[0m \u001b[31m3.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hDownloading optimistix-0.0.10-py3-none-any.whl (84 kB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m84.3/84.3 kB\u001b[0m \u001b[31m3.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hInstalling collected packages: typeguard, jmp, jaxtyping, jaxlib, dm-haiku, jax, equinox, lineax, optimistix, diffrax\n",
            "  Attempting uninstall: typeguard\n",
            "    Found existing installation: typeguard 4.4.1\n",
            "    Uninstalling typeguard-4.4.1:\n",
            "      Successfully uninstalled typeguard-4.4.1\n",
            "  Attempting uninstall: jaxlib\n",
            "    Found existing installation: jaxlib 0.4.33\n",
            "    Uninstalling jaxlib-0.4.33:\n",
            "      Successfully uninstalled jaxlib-0.4.33\n",
            "  Attempting uninstall: jax\n",
            "    Found existing installation: jax 0.4.33\n",
            "    Uninstalling jax-0.4.33:\n",
            "      Successfully uninstalled jax-0.4.33\n",
            "\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n",
            "inflect 7.4.0 requires typeguard>=4.0.1, but you have typeguard 2.13.3 which is incompatible.\u001b[0m\u001b[31m\n",
            "\u001b[0mSuccessfully installed diffrax-0.6.2 dm-haiku-0.0.13 equinox-0.11.11 jax-0.4.38 jaxlib-0.4.38 jaxtyping-0.2.36 jmp-0.0.4 lineax-0.0.7 optimistix-0.0.10 typeguard-2.13.3\n"
          ]
        }
      ],
      "source": [
-        "!pip install diffrax jax jaxlib scipy optax dm-haiku"
+        "try:\n",
        "    import google.colab  # only to ensure that we are inside colab\n",
        "    %!pip install diffrax jax jaxlib scipy optax dm-haiku\n",
        "except ImportError:\n",
        "    print(\"This notebook is running locally, please make sure the packages above are installed\")\n",
        "    pass"
      ]
    },
    {
@@ -150,7 +107,7 @@
        "id": "27e7d7cd-5416-4fa3-8fdb-cfa7efbe35c2"
      },
      "source": [
-        "We use **os**, **random**, **math**, **functools**, **numpy**, **scipy**, **tqdm** and **typing** as helper libraries"
+        "In addition to several widely used numpy and matplotlib libraries:"
      ]
    },
    {
@@ -170,28 +127,8 @@
        "import numpy as np\n",
        "from tqdm import tqdm\n",
        "from scipy.interpolate import griddata\n",
-        "from typing import Any, Callable, Iterable, List, Optional, Tuple, Union"
+        "from typing import Any, Callable, Iterable, List, Optional, Tuple, Union\n",
-      ]
+        "\n",
    },
    {
      "cell_type": "markdown",
      "id": "f3680fcb-71d2-46fb-aac9-5259d1022b0f",
      "metadata": {
        "id": "f3680fcb-71d2-46fb-aac9-5259d1022b0f"
      },
      "source": [
        "For visualizations, we use **matplotlib==3.5.1** and **seaborn**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "id": "6efeb2fb-84b5-4e09-a62b-bbccf327c5c0",
      "metadata": {
        "id": "6efeb2fb-84b5-4e09-a62b-bbccf327c5c0"
      },
      "outputs": [],
      "source": [
        "from matplotlib import pyplot as plt\n",
        "\n",
        "from mpl_toolkits.axes_grid1 import make_axes_locatable\n",
@@ -207,13 +144,7 @@
      "source": [
        "## Physical System SDE\n",
        "\n",
-        "As physical system, we consider the SDE defined by\n",
+        "Next we set up a lambda function for the SDE above with constants $\\lambda = 7$ and $g = 0.03$. "
        "\n",
        "$dx = - \\lambda$ sign$(x) x^2 dt + g dW$\n",
        "\n",
        "with $\\lambda = 7$ and $g = 0.03$.\n",
        "\n",
        "Paths of this SDE will start in either $-1$ or $1$ at $t=0$ and approach $0$ as $t$ increases."
      ]
    },
    {
@@ -277,71 +208,7 @@
      },
      "source": [
        "## Visualization of SDE paths\n",
-        "We first set a *seed* for the SDE paths"
+        "We first set a *seed* for the SDE paths, and plot some example paths from the SDE."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "id": "7592d2ce-c190-4cc8-a382-c6eb1d0d2a00",
      "metadata": {
        "ExecuteTime": {
          "end_time": "2023-11-06T11:31:04.115212500Z",
          "start_time": "2023-11-06T11:31:04.072356600Z"
        },
        "id": "7592d2ce-c190-4cc8-a382-c6eb1d0d2a00"
      },
      "outputs": [],
      "source": [
        "seed = 2022"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "0a33ec0f-098c-4ddd-b01b-7e9b45d7f88d",
      "metadata": {
        "id": "0a33ec0f-098c-4ddd-b01b-7e9b45d7f88d"
      },
      "source": [
        "Next, we define some options for plotting"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "id": "5a56d9d9-a9c3-49fe-a9cf-c098750d66d5",
      "metadata": {
        "id": "5a56d9d9-a9c3-49fe-a9cf-c098750d66d5"
      },
      "outputs": [],
      "source": [
        "#sns.set_context(\"paper\", rc={\"font.size\":8,\"axes.titlesize\":8,\"axes.labelsize\":5})\n",
        "#sns.set(font_scale=1.5, rc={'text.usetex' : True})\n",
        "#sns.set(rc={'axes.facecolor':'black', 'figure.facecolor':'white', 'axes.labelcolor' : \"black\", 'text.color' : \"black\"})"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "id": "qpNPBowBAfPg",
      "metadata": {
        "id": "qpNPBowBAfPg"
      },
      "outputs": [],
      "source": [
        "# NT_DEBUG\n",
        "#print(plt.style.available)\n",
        "#sns_white = \"seaborn-v0_8-white\""
      ]
    },
    {
      "cell_type": "markdown",
      "id": "d2273628-8052-466a-a3ab-ed00cf20795d",
      "metadata": {
        "id": "d2273628-8052-466a-a3ab-ed00cf20795d"
      },
      "source": [
        "Now, we plot paths from for the SDE"
      ]
    },
    {
@@ -369,9 +236,8 @@
        }
      ],
      "source": [
-        "#with plt.style.context(\"seaborn-white\"):\n",
+        "seed = 2022\n",
        "\n",
        "#with plt.style.context(sns_white):\n",
        "if 1:\n",
        "    # paths starting in 1.0\n",
        "    value_one = [r_process(1.0, g, seed+n) for n in range(0, 3)]\n",
@@ -411,16 +277,10 @@
        "id": "dcc8b7c2-88ff-487f-ad05-b1eb1c1056a8"
      },
      "source": [
-        "## Generate Training Data"
+        "As promised, the trajectories start on either side (blue and red), and approach the zero line in a stochastic fashion.\n",
-      ]
+        "\n",
-    },
+        "## Generate Training Data\n",
-    {
+        "\n",
      "cell_type": "markdown",
      "id": "8dce3c11-f5ca-4f0b-988d-5c20a87c522e",
      "metadata": {
        "id": "8dce3c11-f5ca-4f0b-988d-5c20a87c522e"
      },
      "source": [
        "We generate a data set of 250 paths from $t=0$ until $t=10$ with $\\Delta t = 0.02$"
      ]
    },
@@ -586,16 +446,6 @@
        "    return iterate(X, batch_size, shuffle)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "c51e71d3-4217-4dd0-861d-0a4938f47829",
      "metadata": {
        "id": "c51e71d3-4217-4dd0-861d-0a4938f47829"
      },
      "source": [
        "## Neural Network Setup"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "eb481328-018f-408d-8f65-7d5618ddb77f",
@@ -603,7 +453,11 @@
        "id": "eb481328-018f-408d-8f65-7d5618ddb77f"
      },
      "source": [
-        "First, we define the neural network architecture using haiku"
+        "### Neural Network Setup\n",
        "\n",
        "This example employs a neural network $s_\\theta(x,t)$ parameterized by $\\theta$ to approximate the score. It's a simple multilayer perceptron with elu activations and five hidden layers with $30$, $30$, $25$, $20$, and then $10$ neurons for the last hidden layer.\n",
        "\n",
        "The neural network architecture is realized with `haiku`, using a stack of MLPs:"
      ]
    },
    {
@@ -632,7 +486,7 @@
        "id": "9a98e325-e59c-4aec-92de-526e20005bb5"
      },
      "source": [
-        "Then, we initialize the parameters"
+        "The next cell initializes the parameters. The `forward_fn` will be our main handle to evaluate the network `f` later on."
      ]
    },
    {
@@ -659,7 +513,7 @@
        "id": "3f78a2cc-b7a9-4d0e-9ea9-3c897056d5fc"
      },
      "source": [
-        "We define a function to evaluate the model and a function that implements the backpropagation and updating of parameters given the optimizer and model loss"
+        "We also define a function to evaluate the model and a function that implements the backpropagation and updating of parameters given the optimizer and model loss."
      ]
    },
    {
@@ -672,9 +526,6 @@
      "outputs": [],
      "source": [
        "def create_eval_fn(forward_fn, params):\n",
        "    \"\"\"\n",
        "    Evaluate model\n",
        "    \"\"\"\n",
        "    @jax.jit\n",
        "    def eval_model(t, x, rng=None):\n",
        "\n",
@@ -684,9 +535,6 @@
        "\n",
        "def create_default_update_fn(optimizer: optax.GradientTransformation,\n",
        "                             model_loss: Callable):\n",
        "    \"\"\"\n",
        "    Update parameters and backpropagation\n",
        "    \"\"\"\n",
        "\n",
        "    @jax.jit\n",
        "    def update(params, opt_state, batch, rng) -> Tuple[hk.Params, optax.OptState, jnp.ndarray]:\n",
@@ -705,8 +553,8 @@
        "id": "4c92e77d-3f0c-4696-867c-272b787abdb5"
      },
      "source": [
-        "Finally, we implement the ODE solver for the probability flow ODE using Euler steps, </br>\n",
+        "Finally, we implement the ODE solver for the probability flow ODE using Euler steps. \n",
-        "which given an initial state and ground truth (**x_train**) as well as time discretization (**t_train**) compute the L2 loss"
+        "Given an initial state and ground truth (**x_train**) as well as time discretization (**t_train**), it computes the L2 loss:"
      ]
    },
    {
@@ -722,16 +570,12 @@
        "\n",
        "    @jax.jit\n",
        "    def model_loss(model_weights, rng, x_train, t_train):\n",
        "\n",
        "        x = x_train[:,0]\n",
        "        i = 1\n",
        "\n",
        "        loss = 0.0\n",
        "\n",
        "        for t1, t0 in zip(t_train, t_train[1:]):\n",
        "\n",
        "            delta_t = t1-t0\n",
        "\n",
        "            physics_update = physics_operator(x)\n",
        "\n",
        "            # note that we absorb g**2 (constant) in the definition of forward_fn here\n",
@@ -739,27 +583,14 @@
        "                                              jnp.repeat(jnp.tile(t1, 1)[None], x.shape[0], axis=0))[:,0]\n",
        "\n",
        "            x = x - delta_t * (physics_update + score_update)\n",
        "\n",
        "            x_true = x_train[:,i]\n",
        "\n",
        "            loss += jnp.mean(jnp.square(x - x_true))\n",
        "            i += 1\n",
        "\n",
        "        return loss\n",
        "\n",
        "    return model_loss"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "68833303-390b-489a-9372-ffeeba452f6c",
      "metadata": {
        "id": "68833303-390b-489a-9372-ffeeba452f6c"
      },
      "source": [
        "## Visualization of Score Field"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "47c38a2c-0c63-49b9-9172-b5c2b369082d",
@@ -767,8 +598,10 @@
        "id": "47c38a2c-0c63-49b9-9172-b5c2b369082d"
      },
      "source": [
-        "We implement the visualization of the score field, which we will be using during and after training. </br>\n",
+        "### Visualizing the Score Field\n",
-        "First, we set up the grids used for plotting."
+        "\n",
        "While the problem itself is very clear, it's actually much less obvious what the score function underneath should look like. Hence, we'll show different version of score functions below, and the next cell implements a helper function that takes care of the visualization of the score field. \n",
        "We'll sample the score function on a regular grid to show it as an image, and hence first set up the corresponding grids."
      ]
    },
    {
@@ -795,7 +628,7 @@
        "id": "1dced7ca-2152-4805-8f49-6ea632b8fc57"
      },
      "source": [
-        "This function plots the score field"
+        "The following function `save_snapshot_score` plots a score function as an image."
      ]
    },
    {
@@ -814,12 +647,9 @@
        "\n",
        "def save_snapshot_score(params, forward_fn, meshgrid, step=None, savename=None):\n",
        "    eval_fn = create_eval_fn(forward_fn, params)\n",
        "\n",
        "    results = None\n",
        "\n",
        "    X, Y = meshgrid\n",
        "    full_domain = np.hstack((X.flatten()[:,None], Y.flatten()[:,None]))\n",
        "\n",
        "    full_domain_scaled = full_domain * scaling_factor\n",
        "\n",
        "    generator = iterbatches_grid(full_domain_scaled, 1000, shuffle=False)\n",
@@ -849,7 +679,6 @@
        "    fig.set_dpi(dpi)\n",
        "\n",
        "    ax = axes\n",
        "\n",
        "    u_pred = final_results[0][:].flatten()\n",
        "    U_pred = griddata(full_domain_scaled, u_pred.flatten(), meshgrid_scaled, method='cubic')\n",
        "\n",
@@ -857,30 +686,40 @@
        "    h = ax.imshow(jnp.flip(U_pred / (g**2), axis=0), cmap='jet',\n",
        "                  extent=[ x_scaled.min(), x_scaled.max(), y_scaled.min(), y_scaled.max()],\n",
        "                  aspect='auto', vmin = -vmax, vmax = vmax)\n",
        "    # divider = make_axes_locatable(ax)\n",
        "    # cax = divider.append_axes(\"right\", size=\"5%\", pad=0.10)\n",
        "    cbar = fig.colorbar(h)\n",
        "    # cbar.ax.tick_params(labelsize=15)\n",
        "    ax.set_xlabel(\"t\")\n",
        "    ax.set_ylabel(\"x\")\n",
        "\n",
-        "    # ax.tick_params(labelsize=15)\n",
+        "    plt.figure(figsize=(10, 6)) # NT_DEBUG , test\n",
        "\n",
        "    # axTransform = {'horizontalalignment':'right',\n",
        "    #          'verticalalignment':'center', 'transform':ax.transAxes}\n",
        "\n",
        "    # ax.grid(False)\n",
        "\n",
        "    # ax.patch.set_edgecolor('black')\n",
        "    # ax.patch.set_linewidth('1')\n",
        "\n",
        "    #fig.canvas.draw()\n",
        "    #image = np.frombuffer(fig.canvas.tostring_rgb(), dtype='uint8')\n",
        "\n",
        "    if savename:\n",
        "        plt.savefig(f'{savename}.svg', transparent=True)\n",
-        "\n",
+        "    plt.show()"
-        "    plt.show()\n"
+      ]
    },
    {
      "cell_type": "markdown",
      "id": "d7f5ca77",
      "metadata": {},
      "source": [
        "Let's take a loot at the untrained score function `forward_fn`:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "0a7c7693",
      "metadata": {},
      "outputs": [],
      "source": [
        "save_snapshot_score(params, forward_fn, meshgrid)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "5167f0d5",
      "metadata": {},
      "source": [
        "Not too surprisingly, it contains random but smooth transitions. Let's revisit it after a first training run."
      ]
    },
    {
@@ -892,8 +731,9 @@
      "source": [
        "## Training and Sliding Window Method\n",
        "\n",
-        "We implement the sliding window method and training algorithm below. </br>\n",
+        "Next we implement functions for network training. We implement the _sliding window_ method and corresponding training algorithm below. \n",
-        "The sliding window method starts with window size **ROLLOUT_start**, which is increased by **ROLLOUT_add** every **steps** epochs for **ROLLOUT_increases** times. </br>\n",
+        "\n",
        "The sliding window method starts with window size **ROLLOUT_start**, which is increased by **ROLLOUT_add** every **steps** epochs for **ROLLOUT_increases** times. \n",
        "To reduce the number of points on a trajectory, we use subsampling (defined by **subsample**)"
      ]
    },
@@ -909,9 +749,7 @@
        "def update_network(params_, forward_fn, dataset, rng, steps = 20, ROLLOUT_increases = 18, max_training_time=None,\n",
        "                   ROLLOUT_start = 4, ROLLOUT_add = 2, lr=5e-4, bidirectional=False, subsample=5):\n",
        "\n",
-        "    # Learning rate\n",
+        "    scheduler = optax.piecewise_constant_schedule( init_value=lr)\n",
        "    scheduler = optax.piecewise_constant_schedule(\n",
        "        init_value=lr)\n",
        "\n",
        "    # Optax otimizer using Adam\n",
        "    opt = optax.chain(\n",
@@ -926,15 +764,12 @@
        "\n",
        "    avg_loss = 0\n",
        "    grad_updates = 0\n",
        "\n",
        "    ROLLOUT = ROLLOUT_start\n",
        "    history = []\n",
        "\n",
        "    # Iterate through all sliding window sizes\n",
        "    for _ in range(ROLLOUT_increases+1):\n",
        "\n",
        "        print('Window size: ', ROLLOUT)\n",
        "\n",
        "        pbar = tqdm(range(steps))\n",
        "\n",
        "        # Iterate through all epochs\n",
@@ -952,21 +787,17 @@
        "                    x_train = x_train[:, -max_training_time:]\n",
        "                    t_train = t_train[-max_training_time:]\n",
        "\n",
        "\n",
        "                # Use subsampling to reduce number of points on trajectory\n",
        "                x_train = x_train[:, ::subsample]\n",
        "                t_train = t_train[::subsample]\n",
        "\n",
        "\n",
        "                # Iterate through trajectory\n",
        "\n",
        "                for t in range(x_train.shape[1]):\n",
        "\n",
        "                    # Select values based on position and window size\n",
        "                    x_train_sub = x_train[:, t:t+ROLLOUT]\n",
        "                    t_train_sub = t_train[t:t+ROLLOUT]\n",
        "\n",
        "\n",
        "                    # Compute loss for current part of trajectory\n",
        "                    params_, opt_state, batch_loss_backward = grad_update(params_, opt_state, [x_train_sub, t_train_sub], rng=rng)\n",
        "\n",
@@ -991,44 +822,6 @@
        "    return params_, history"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "6d8f989f-87dc-4b1a-87d6-0f7ad4536115",
      "metadata": {
        "id": "6d8f989f-87dc-4b1a-87d6-0f7ad4536115"
      },
      "source": [
        "We plot the score field before any training is done"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 24,
      "id": "e5b85282-57e9-431f-941a-96f772188569",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 760
        },
        "id": "e5b85282-57e9-431f-941a-96f772188569",
        "outputId": "aa10e282-77cf-4f3c-cc2d-e5b38c41a5f6"
      },
      "outputs": [
        {
          "data": {
            "image/png": "iVBORw0KGgoAAAANSUhEUgAABCkAAALmCAYAAABvmM48AAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjguMCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy81sbWrAAAACXBIWXMAAB7CAAAewgFu0HU+AAD8nUlEQVR4nOzde7wkdX3n//c5VWcuDChe0ICDosIIxrhhFYJBxNu6uygSTKIkMUZFUVdZyRrQRKPurkQRFZTHGi+gxKhAYhZIRB7rIwliICBO5OclQBCFCCNGSTTc5nKqTv/+6K7qb33q862qPtPd0+ec1/PxaLu66ltVfYaZeVif+Vzmer1eTwAAAAAAAHvY/J7+AgAAAAAAABJBCgAAAAAAMCMIUgAAAAAAgJlAkAIAAAAAAMwEghQAAAAAAGAmEKQAAAAAAAAzgSAFAAAAAACYCQQpAAAAAADATCBIAQAAAAAAZgJBCgAAAAAAMBMIUgAAAAAAgJlAkAIAAAAAAMwEghQAAAAAAGAmEKQAAAAAAAAzgSAFAAAAAACYCQQpAAAAAADATEj39BfAyrZjxw59+9vfliTtt99+SlN+SwEAAAB7SpZl+slPfiJJ+oVf+AVt2LBhD3+j7rIs049+9KM9+h1+7ud+jmeaPYxffeyWb3/72zryyCP39NcAAAAAYNxwww064ogj9vTX6OxHP/qRDjzwwD36He68805t3rx5j36HtY5yDwAAAAAAMBPIpMBu2W+//crti244QL9x5Mv24LdZzVJJe0vaIGlfSRslbTKf95W00F++XtJegyXrg/eHDC4VvifOuvVmX7lmUemmHUqSTBs37dB8sqSN2q4F7dJ67dKCdmmdFrVOuzSvXBu1Q4kybdQOzau/NlGuddo1WLNUbq/TLiXKtVHby/3zyrV+sK44d0G7lCrXghaVKNM6LWphcO46Lfbf811KslwLO3tKc2kul7RLUha875S0JGnH4Jd4l6R8sD83a5eCc3JzXMG6PFi/IzgvD953Bud45+bmOpnUy6Usl7IseF/qL1kcXKb4MfLBvth7Nngp2M7M/h3mOAAAWDnuk3T+YDv8/+orzWvV/3+/03C/pE9O6V5oR5ACuyWs19pv/+KpF+O3IGkf9YMRD1c/AvHwwed9gs8Lw+WJ+n/CFwaH08HyDRrGN/Y273ZfGBvZuyelmeb23q6FDTs1v2GX1iU7lWqXFrRzEKDYqfWDgMN67ay8J8q1l9YrUV7u669PlWh9sH5+EHBIlSjTXpISZVo/OHedpPXapUTSei0pUap1ygefcyXqaX0+pySb07odgyDFDvWf0ov3LNgO36V+ECEL3sPtjZG16801s8Gvf26uIUlzg/1zg/3r1M9pywb75gfvGqzrSb2elPWkxfl+cGJxbhigKAIP4bskbTfvsfVy9q035wIAgJVpJfdW2FvSQ/f0l8AesXJ/12LmJMr39FdYxYpoQ/FKNfzjaz6nwSmp8zm2xr5qxzMpzZWk/UfbJMmVKlei+nuirHyXVNuXRNYP1/b3pZX34Tn9dVmw1lwry5VkS/0ARZilIA2DBrl5D49lZl0hXGvPD1MO7P7M7I+tyevbvUH2xGLWz6BYLD7LDzjYdwVrugQoFs25AAAAe0L4f0mncS/MDnpSACuK/Su0CF6YXbFAg5z32OXteYEkDQMQ1QBD/9ThexhwCPdVb5tXghqxdWHQw5MoU5r3AxSJ99TuBRTkHLeXt+fkqkcBbMAjPGav0RYFGNy/N1i3aN5jFp33zOyz67OGYwAAAMC0ETTC2JBJMWlNf1wX/N1eTMN772g+zZWmw//O4X/ztgBCuKbY9tYmbnAidwMh1YyL+rXmmoICNsvBEwtGhMdjwrXeutysswGUwXs2WJcFn21WhNQcVCiOxQIS9hgBCgAAsKc5/xQ30XthdvDfA1gRFoJ3+yrSHQaNDEYp4WhbX3nl/QyKwctmRqQmkOCVYoTH6mUbmerlIsN14b388pB8kEWRK+lXplSzIpxSCrf0wgsWNJWF2PNjJSINgYiahgCFVA8ihJfzyjdsCUesvIMyDwAAAOxplHtgbObJpJiQBfO+jAo9L3OiaxBDktKe5sMMiqSa1RAKy0DCNTYwYffFyjvCHhXDr169TquuGRWxko7Ytr1mrBdFcZ79LrFeFIqXetisiFECDrEmm+F1yaIAAADAnkQmBbAixBpIRP4IN/WgaPpT75WHRNZXMxnqAQnv3W6PY1/lPm0NMyW/b0SorQSkEGuo6V3Pbo8Yz2vrRVFb72x3KfMAAACYFTTOXLv47wHMPJviUJR5FILPtQyIyCW6vsrrZIMyj6zWNLMQn8gxXOv1kxhlqoctBfGyM2piDTObyj+kaqZDU5lI7JqxdcuY6CH5vShi71I14yL8kWJTPAhWAAAAYBYQpMDY0Dhzkry2Qd44UsWDDaNkV1TO6Q13p3ml1CPWjyJ8t/ts+UfbVA8bzAiDEpXmmcFUjzmvZ4SCbS9IEB63vSrsuU3bBa9HhZdVYYIXsQCF5E/v8GTOOi8wYY8BAADMChpnrl30pABWlJa/qkf5G7YtYFHJpOhnUQxP7R6QSk1wQYqXcoT7YxM7Yud30nRaU68Jy+s1UexvKxdp6lPh8HpR2HevF0WXaR8AAADArCFohLFJtLSnv8IqVWRL2MaZYenH3HD5KKUcTWUg5f6sHD1qSz28SRxehkTYs8I7vz7pozratF7qYUo+YlM9vFIPuy9WsiE1l4R4pRqhWDmHnO3Bu82iWCw+q7lJZsEr9YiVg3j7AAAAZgU9KdYuMimAmRZmToRBCduXQn7wIdzfKWPCOy8IGKTVQER/md9fIjwWZkXY84v1thGnDVaE5xbb/fVZvGGmF3CIBRa8HhRyriNzPLx27LreGntsmdpGjnpJGrEgBwAAALCnEaQAZl5TRZ4TYx5lisduho2bRpD66zPzuRqs8Hh9K5ru7+paftF2bLklISNGA5bbi8KyPSma1gAAAACzgMwWjA2NMych1tEysrTrNI8uawZNM+eLMo9B08ywzEIaZkHY5pnea3gsq/x+aZrqMfzxvOkeDQ0zvawFm+3QJcvCmwxiyzXarmvXtUz0kEbrReFlUXgBjVjpBwAAwKyhcebaRSYFxoYgxaTEelI4/SiKw+GypmMdR5AWo0fLz6bXRLjf6x/RP1YvFfH6WQyvVS/7sKUeNbZEI9z2+kBIfmlIsSbWxyI2/cPrYSFn21xzOb0omnpNeL8cTb0sAAAAgFlB0AhYURriyU3dhWL9KtqY4ETIZkY0iU3zGGW9d19J1YaZkt8PQs6+2Fe2gYeY3SjnmFSUIMyMaPsxyaIAAACzjMaZaxf/PYCZFmZRLAT7iv3OH+GGjIjOZSCDUo/ytLLUI6s1zSzYqR5h6Ub/eHUyR3G+VzZSrG/KypCkNDfRBFuiEW7bTAkpXhoic05sGsgI5RyVNcF3HVcWRay8Q856AAAAYFZR7oGxmafcYwK8eo2w3MM5ZJfFghBqWS9JaV6WeoSBBylequG92332fO+a4TpvOog71UOK1zsU717wwSvZCNeH17O9KOx+u9YGL7zvZTT1oWg4LVreEV7TbgMAAACzhEwKYMXw/rg65R9dpnqE29ESkfi/ubf1H/EyLWzWRds1bY+KLvfdLTYg0fVWTSUlTesH714WhRSf6OFlR3T9ijTLBAAAKwWNM9cuMikwNrXUe+ymsLwjNfuMtmyJLpkUzvt8mitNqxkMqZfR0NDksr89zJTwzq/v86d6FNftPNXDK/XwsihstoTNkGgr4/CaatrrOc0zey2ZFF1Hj7aVg0jVHxEAAACYVQSNgJnWsSfFqD0o7DG7RhqUegxfUn20aP+0aslH2HPCloN4Y0SH66pBiNYxp2HDzLb+EU2vQtfAhdTcZ8ILSqi+pqkXRbjce4+NHG0r7yCLAgAArBQ0zly7yKQAZpb312XHv0Lbpnm0lnr0IgeqmqdwdPu3+1HLQKQxZ+00fc1YoMH73NYoou1eDaeOmg0RriMwAQAAgJWEIAXGJsnG+OCIgVhaRJFNMRfPnrCnFO8dm2bOh00zE7+0Q4pnPgyP5bVjtt9E01SP6L6wYaaXRWGbYrZlSnjXsMcVHDOlG5Vj9ho280LxLIpYRoTXi6KtrINeFAAAAFhpyGwBZpJX3hFum1IPyQ82dOlT4b6yssRDUiUo0T89d/fb0aN+z4owwOH3s6heq55pUQmIdelF0dafwiv/2J1eFN53Gawv+lAsBu/LHTlqgw+xkaMEKAAAwEoT/j/gadwLs4NMCozNfN6tRACjWIhsG23hxlimRdtpgyyK4anV3hHV/e2ZNLHsiepX9ad6lD0w8iJQEYwdXY7Y17XBihhvjQ1WtNwv253vb25L3wkAAACsBmRSADOrqUYjsrytUWZbJoXU70cxKPUYXnoYnJCGwQYbZPAaZNpxpLFSj7apHv21w1IPl9fM0m7b9V75R3gs1jCzy/m27CMo9SiXBSNHm5plhu+xQESsuSYAAMBKwwjStYv/HhibZEz/KgzJ7z8RjiEd9KMolnQp8+gcrMg0X0zzSIelGHYEadvoUb9nRdZxX/XVeaqH1DhRwy3viI0nVcM59rxYLwpV13gTPaR4yUbTpI4uZSEAAADASkOQAlhROsSTm0o7Rij7SNNY3UJ8qocNUNgsiti5XUpFpIapHm0NM8N3yS/JaJvkEdNU3pHF19gARZeSDTuetA1ZFAAAYKViBOnaRU8KYCYtqN4uqNg2f2V72ROhtkwKu1Yqm2YmaV4LMIQBCFva0V9T7VnhNdy059t9I0/1iPECFDb7wZ7fpaSjaTqIVyKiasPMohdFWOoR+8qxUo+mLAp7TQAAAGClIGiEsWn4h3csS1je4ZV/yA9MtJV8NI4g7ZX/IaujR+2kjmp5hy0HaRtF6pWKeIGJyvXzvCz1mPPKL2wWRWyfN5Gj6XqSX9rhBUliZLine truncated
            "text/plain": [
              "<Figure size 1200x800 with 2 Axes>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "save_snapshot_score(params, forward_fn, meshgrid)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "0815459e-63ac-47df-bf99-8c00ad984a5d",
@@ -1036,9 +829,10 @@
        "id": "0815459e-63ac-47df-bf99-8c00ad984a5d"
      },
      "source": [
-        "### Single step loss - Sliding window size 2\n",
+        "### Single step loss \n",
-        "We train the score network with the single step loss for 2500 epochs using the probability flow ODE (see Single steps, Section 3) </br>\n",
+        "We train the score network with the single step loss for 2500 epochs using the probability flow ODE (see Single steps, Section 3).\n",
-        "First, we use subsampling with factor 5 to reduce the number of points on the trajectory"
+        "Single step is realized here with a _sliding window_ of size 2, i.e. going one step forward in time for each sample. Increases of the unrolling are disabled via `ROLLOUT_increases=0` by default below. This speeds up training for default runs, but feel free to enable the increases in subsequent tests.\n",
        "We also use subsampling with factor 5 to reduce the number of points on the trajectory. After training, the resulting score function is visualized again."
      ]
    },
    {
@@ -1080,9 +874,11 @@
        }
      ],
      "source": [
        "ROLLOUT_increases=0\n",
        "\n",
        "key = jax.random.PRNGKey(seed)\n",
        "params_single_step, history = update_network(params, forward_fn, data, key, steps = 2500,\n",
-        "                                             ROLLOUT_start=2, ROLLOUT_increases=0, bidirectional=True, lr=1e-3,\n",
+        "                                             ROLLOUT_start=2, ROLLOUT_increases=ROLLOUT_increases, bidirectional=True, lr=1e-3,\n",
        "                                             subsample=5)\n",
        "\n",
        "\n",
@@ -1096,7 +892,9 @@
        "id": "a6275c68-6492-4ba0-a3c5-40b10055aaf3"
      },
      "source": [
-        "Plot the training loss"
+        "This should look more meaningful than the initial random state! The two bands of trajectories from -1 and 1 should be visible. E.g., the blue region at the top indicates a large negative gradient, driving samples ending up in this  region down by a large distance, towards the correct trajectory in the positive region (as shown further above). Note that some parts, typically closer to the end time ($t=10$ on the right), can show sign changes and mostly random content again. This is caused by the network primarily being trained in the vicinity of samples it has seen in the training data. Far away from it, it has \"not learned\" how to properly transform samples into admissible ones.\n",
        "\n",
        "We also plot the training loss below, to show the training progressing for this first training step."
      ]
    },
    {
@@ -1138,9 +936,8 @@
        "id": "e0b02cf8-1b99-418a-a9c9-e440972fd5e4"
      },
      "source": [
-        "We decrease the learning rate and train without subsampling of the trajectories\n",
+        "For the next training step, the goal is to improve the quality of the learned score:\n",
-        "\n",
+        "We decrease the learning rate and train without reducing the trajectories via \"subsampling\", meaning all samples in the training data are provided to the network."
        "[NT_DEBUG , benjamin - why no subsampling?]"
      ]
    },
    {
@@ -1183,12 +980,22 @@
      ],
      "source": [
        "params_single_step, history = update_network(params_single_step, forward_fn, data, key, steps = 1000,\n",
-        "                                             ROLLOUT_start=2, ROLLOUT_increases=0, bidirectional=True, lr=1e-4,\n",
+        "                                             ROLLOUT_start=2, ROLLOUT_increases=ROLLOUT_increases, bidirectional=True, lr=1e-4,\n",
        "                                             subsample=1)\n",
        "\n",
        "save_snapshot_score(params_single_step, forward_fn, meshgrid)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "2e75b45b",
      "metadata": {},
      "source": [
        "This typically results in an even _cleaner_ score function, containing fewer sign changes outside the two bands of the ground truth trajectories. The runtime above was reduced via the `ROLLOUT_increases=0` setting. Once you enable it, e.g., with `ROLLOUT_increases=10` potentially with an increased number of iterations, you should see further improvements of the score field in the visualization.\n",
        "\n",
        "The training in this stage is also more stable, as shown in the corresponding plot below."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 28,
@@ -1228,9 +1035,10 @@
        "id": "82b8cba3-feaf-4f27-b729-e3b0c6f0435b"
      },
      "source": [
-        "## Single-step Inference\n",
+        "## Sampling SDE and ODE trajectories\n",
-        "We use the trained score with diffrax to solve the probability flow ODE and simulate paths from the reverse-time SDE. </br>\n",
+        "\n",
-        "For this task, it is sufficient to train with the single step objective only and not consider higher sliding window sizes."
+        "Above we've only looked at the learned score function, and thus we'll evaluate the trajectories produces by integrating the two DE variants.\n",
        "We use the trained score with `diffrax` to solve the probability flow ODE and simulate paths from the reverse-time SDE. "
      ]
    },
    {
@@ -1241,7 +1049,7 @@
      },
      "source": [
        "### Reverse-time SDE\n",
-        "We define a function to simulate paths from the reverse-time SDE with diffrax"
+        "We define a function to simulate paths from the reverse-time SDE with `diffrax`, where `WeaklyDiagonalControlTerm(diffusion, brownian_motion)` takes care of the noise:"
      ]
    },
    {
@@ -1257,7 +1065,6 @@
        "    key = jr.PRNGKey(seed)\n",
        "    initial_shape = (1,)\n",
        "    y1 = jnp.ones(shape=initial_shape) * initial_value\n",
        "\n",
        "    t0, t1 = 0.0, 10.0\n",
        "    dt0 = 0.01\n",
        "\n",
@@ -1270,26 +1077,18 @@
        "    terms = MultiTerm(ODETerm(drift), WeaklyDiagonalControlTerm(diffusion, brownian_motion))\n",
        "\n",
        "    solver = dfx.Euler()\n",
        "\n",
        "    t0 = jnp.array(0.0)\n",
        "\n",
        "    args = None\n",
        "    tprev = jnp.array(t1)\n",
        "    tnext = jnp.array(t1 - dt0)\n",
        "    y = y1\n",
        "\n",
        "    state = solver.init(terms, tprev, tnext, y1, args)\n",
        "\n",
        "    y_list = []\n",
        "\n",
        "    for i in range(((t1-t0) / dt0).astype(int)):\n",
        "\n",
        "        y, _, _, state, _ = solver.step(terms, tprev, tnext, y, args, state, made_jump=False)\n",
        "\n",
        "        tprev = tnext\n",
        "\n",
        "        tnext = jnp.array(jnp.maximum(tprev - dt0, t0))\n",
        "\n",
        "        y_list.append(y)\n",
        "\n",
        "    return y_list"
@@ -1302,7 +1101,7 @@
        "id": "4a218961-65ea-4e8f-a83f-59f234282910"
      },
      "source": [
-        "Next, we plot ten different paths from the reverse-time SDE"
+        "and plot ten different paths sampled from the reverse-time SDE:"
      ]
    },
    {
@@ -1330,10 +1129,8 @@
        }
      ],
      "source": [
        "#with plt.style.context(\"seaborn-white\"):\n",
        "if 1:\n",
        "    combinations = zip(np.linspace(-0.065,0.065, 6), range(6))\n",
        "\n",
        "    fig, ax = plt.subplots(figsize=(6,4))\n",
        "\n",
        "    for init, seed_ in combinations:\n",
@@ -1353,10 +1150,19 @@
        "\n",
        "    plt.xlim([10.25, -0.25])\n",
        "    plt.ylim([-1.25,1.25])\n",
        "\n",
        "    plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "9abbff53",
      "metadata": {},
      "source": [
        "Note that the trajectories are shown along integration of the reverse-time SDE, and thus they appear mirrored along y in comparison to the ground truth ones above. With default settings, these trajectories should look quite good: exhibiting noisy motions during the initial phase, and reliably converging to $1$ and $-1$.\n",
        "\n",
        "If you reduce accuracy or training duration, you'll see that the score will contain erroneous regions, giving particles that shoot off before reaching $t=0$."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "14bce165-f4e9-4485-81f8-906b6e042538",
@@ -1364,17 +1170,10 @@
        "id": "14bce165-f4e9-4485-81f8-906b6e042538"
      },
      "source": [
-        "### Probability flow ODE"
+        "### Probability flow ODE\n",
-      ]
+        "\n",
-    },
+        "The probability flow ODE, as deterministic counterpart of the SDE, is an interesting variant that is obtained by removing the `brownian_motion` diffusion term.\n",
-    {
+        "Analogously, we define a function to solve the probability flow ODE using only the drift term:"
      "cell_type": "markdown",
      "id": "19b1deb0-6bb1-4174-b9ee-89826a083082",
      "metadata": {
        "id": "19b1deb0-6bb1-4174-b9ee-89826a083082"
      },
      "source": [
        "Analogously, we define a function to solve the probability flow ODE"
      ]
    },
    {
@@ -1390,7 +1189,6 @@
        "    key = jr.PRNGKey(seed)\n",
        "    initial_shape = (1,)\n",
        "    y1 = jnp.ones(shape=initial_shape) * initial_value\n",
        "\n",
        "    t0, t1 = 0.0, 10.0\n",
        "    dt0 = 0.01\n",
        "\n",
@@ -1399,26 +1197,19 @@
        "\n",
        "    terms = ODETerm(drift)\n",
        "    solver = dfx.Euler()\n",
        "\n",
        "    t0 = jnp.array(0.0)\n",
        "\n",
        "    args = None\n",
        "    tprev = jnp.array(t1)\n",
        "    tnext = jnp.array(t1 - dt0)\n",
        "    y = y1\n",
        "\n",
        "    state = solver.init(terms, tprev, tnext, y1, args)\n",
        "\n",
        "    y_list = []\n",
        "\n",
        "    for i in range(((t1-t0) / dt0).astype(int)):\n",
        "\n",
        "        y, _, _, state, _ = solver.step(terms, tprev, tnext, y, args, state, made_jump=False)\n",
        "\n",
        "        tprev = tnext\n",
        "\n",
        "        tnext = jnp.array(jnp.maximum(tprev - dt0, t0))\n",
        "\n",
        "        y_list.append(y)\n",
        "\n",
        "    return y_list"
@@ -1431,7 +1222,7 @@
        "id": "8a7d2f46-dd03-4ed6-80eb-d24f9aa14d1c"
      },
      "source": [
-        "Finally, we plot the solutions to the probability flow ODE"
+        "Finally, we can plot the solutions to the probability flow ODE"
      ]
    },
    {
@@ -1486,14 +1277,17 @@
      ]
    },
    {
-      "cell_type": "code",
+      "cell_type": "markdown",
-      "execution_count": 32,
+      "id": "6be4444a",
-      "id": "A6wd2TRWcqwN",
+      "metadata": {},
-      "metadata": {
+      "source": [
-        "id": "A6wd2TRWcqwN"
+        "They show the influence of the score without random perturbations, and in this case provide mostly horizontal motions with a bifurcation towards the two end points.\n",
-      },
+        "\n",
-      "outputs": [],
+        "## Next steps\n",
-      "source": []
+        "\n",
        "* In this notebook it's worth experimenting with the unrolling steps (cf. `ROLLOUT_increases` above), to evaluate their effect on learning the score.\n",
        "* In addition, it's interesting to intentionally deteriorate the quality of the score with short training runs. This should results in chaotic and erroneous trajectories."
      ]
    }
  ],
  "metadata": {
@@ -1114,7 +1114,12 @@
  journal={arXiv:2209.03003}, year={2022} 
 }
-
+@inproceedings{chung2023dps,
  title={Diffusion posterior sampling for general noisy inverse problems},
  author={Chung, Hyungjin and Kim, Jeongsol and Mccann, Michael and Klasky, Marc and Ye, Jong Chul},
  booktitle={International Conference on Learning Representations},
  year={2023}
 }