<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 This notebook is a short guide that covers issues of interest including
Step1:  Just like how each ordinary differential equation (ODE) is governed by a vector field, a stochastic differential equation (SDE) is governed by two vector fields, which are called the drift and diffusion functions
Step2:  The functions f and g are arbitrarily chosen for demonstration purposes. The attributes noise_type and sde_type must be included to inform the solver about how to interpret the SDE, and determine which numerical methods are available. See below for more detail on how the output of g depends on the different noise types.
We use diagonal here, meaning the output of g should be a vector with the same shape as input y, and it is an element-wise function. 
Note that for any other noise type, we expect the output of g to be a matrix, and a batch matrix-vector product will be performed with the Brownian motion.
The codebase supports both Itô SDEs based on Itô integrals and Stratonovich SDEs based on Stratonovich integrals.
Now we instantiate an object of the SDE class and call the function sdeint on it.
Step3:  For Itô SDEs method='euler' means the strong order 0.5 Euler-Maruyama method is used. Other possible methods include the strong order 1.0 milstein and the strong order 1.5 srk, both of which are of slightly higher order. If method is set to None, an appropriate solver would be chosen based on noise_type and sde_type under the hood.
To solve a Stratonovich SDE with drift and diffusion defined above, we need only change sde_type to stratonovich and supply an appropriate method of choice (e.g. the midpoint method)
Step4:  For Stratonovich SDEs, the methods midpoint, euler_heun, heun, milstein, and log_ode are supported. 
The drift and diffusion functions don't necessarily have to be called f and g. They can be methods with any name, so long as we provide these names to the solver when they differ from the default. The following is an example where h is used as the drift.
Step5:  Trivially, the previous code may be adapted to run on GPUs, just by moving all tensors to a GPU
Step6:  A side note is that multi-GPU data parallel is possible with the existing codebase, but the use case has not been tested out extensively and may require defining non-standard SDE classes and methods. 
2. Explicit control over randomness from the Brownian motion
To gain control over the randomness, we draw Brownian motion samples by instantiating objects of classes BrownianInterval.
Step7:  We can also create the Brownian motion on GPUs by specifying device
Step8:  Having a Brownian motion object helps us gain control over the randomness better. We can feed the object into the solver such that the solver's solution is conditioned on this path.
Step9:  3. Noise type of SDEs affects which solvers can be used and what strong orders can be attained<a id='noise_type'></a>
The supported noise types of this codebase are "diagonal", "additive", "scalar", and "general".
The following is a simple summary of each type
Step10:  Switching to adjoint-mode gradient computation is as simple as replacing sdeint with sdeint_adjoint
<ASSISTANT_TASK:>
Python Code:
import torch
from torch import nn
import os
import sys
module_path = os.path.abspath(os.path.join('..'))
if module_path not in sys.path:
    sys.path.append(module_path)
%matplotlib inline
import matplotlib.pyplot as plt
import torchsde
def plot(ts, samples, xlabel, ylabel, title=''):
    ts = ts.cpu()
    samples = samples.squeeze().t().cpu()
    plt.figure()
    for i, sample in enumerate(samples):
        plt.plot(ts, sample, marker='x', label=f'sample {i}')
    plt.title(title)
    plt.xlabel(xlabel)
    plt.ylabel(ylabel)
    plt.legend()
    plt.show()
Explanation: This notebook is a short guide that covers issues of interest including:
1. Defining and solving SDEs
1. Gaining control over the randomness (path generation conditional on a Brownian motion sample)
1. Different noise types of SDEs
1. Gradient computation/backpropagation
End of explanation
class SDE(nn.Module):
    def __init__(self):
        super().__init__()
        self.theta = nn.Parameter(torch.tensor(0.1), requires_grad=True)  # Scalar parameter.
        self.noise_type = "diagonal"
        self.sde_type = "ito"
    
    def f(self, t, y):
        return torch.sin(t) + self.theta * y
    
    def g(self, t, y):
        return 0.3 * torch.sigmoid(torch.cos(t) * torch.exp(-y))
Explanation: Just like how each ordinary differential equation (ODE) is governed by a vector field, a stochastic differential equation (SDE) is governed by two vector fields, which are called the drift and diffusion functions:
$$dy(t) = \underbrace{f(t, y(t), \theta_f)}{\text{drift}} dt + \underbrace{g(t, y(t), \theta_g)}{\text{diffusion}} dW(t).$$
The output of $f$ is of the same size as the $d$-dimensional state, whereas the output of $g$ may be a matrix of size $(d, m)$. $W(t)$ is the Brownian motion (aka Wiener process), and it is $m$ dimensional. It is a stochastic process, and each random draw produces a function of time. 
Note that the second term on the right hand side $gdW$ is based on matrix-vector products.
1. Solving a simple SDE
To implement an SDE, we create a class with the functions f and g:
End of explanation
batch_size, state_size, t_size = 3, 1, 100
sde = SDE()
ts = torch.linspace(0, 1, t_size)
y0 = torch.full(size=(batch_size, state_size), fill_value=0.1)
with torch.no_grad():
    ys = torchsde.sdeint(sde, y0, ts, method='euler')  # (t_size, batch_size, state_size) = (100, 3, 1).
plot(ts, ys, xlabel='$t$', ylabel='$Y_t$')
Explanation: The functions f and g are arbitrarily chosen for demonstration purposes. The attributes noise_type and sde_type must be included to inform the solver about how to interpret the SDE, and determine which numerical methods are available. See below for more detail on how the output of g depends on the different noise types.
We use diagonal here, meaning the output of g should be a vector with the same shape as input y, and it is an element-wise function. 
Note that for any other noise type, we expect the output of g to be a matrix, and a batch matrix-vector product will be performed with the Brownian motion.
The codebase supports both Itô SDEs based on Itô integrals and Stratonovich SDEs based on Stratonovich integrals.
Now we instantiate an object of the SDE class and call the function sdeint on it.
End of explanation
sde.sde_type = "stratonovich"
with torch.no_grad():
    ys = torchsde.sdeint(sde, y0, ts, method="midpoint")  # (t_size, batch_size, state_size) = (100, 3, 1).
plot(ts, ys, xlabel='$t$', ylabel='$Y_t$')
Explanation: For Itô SDEs method='euler' means the strong order 0.5 Euler-Maruyama method is used. Other possible methods include the strong order 1.0 milstein and the strong order 1.5 srk, both of which are of slightly higher order. If method is set to None, an appropriate solver would be chosen based on noise_type and sde_type under the hood.
To solve a Stratonovich SDE with drift and diffusion defined above, we need only change sde_type to stratonovich and supply an appropriate method of choice (e.g. the midpoint method):
End of explanation
class SDENewName(nn.Module):
    def __init__(self):
        super().__init__()
        self.theta = nn.Parameter(torch.tensor(0.1), requires_grad=False)  # Scalar parameter.
        self.noise_type = "diagonal"
        self.sde_type = "ito"
    
    def h(self, t, y):
        return torch.sin(t) + self.theta * y
    
    def g(self, t, y):
        return 0.3 * torch.sigmoid(torch.cos(t) * torch.exp(-y))
    
sde_new_name = SDENewName()
with torch.no_grad():
    # Supply a dictionary to the argument `names`.
    ys = torchsde.sdeint(sde_new_name, y0, ts, method='euler', names={'drift': 'h'})
plot(ts, ys, xlabel='$t$', ylabel='$Y_t$')
Explanation: For Stratonovich SDEs, the methods midpoint, euler_heun, heun, milstein, and log_ode are supported. 
The drift and diffusion functions don't necessarily have to be called f and g. They can be methods with any name, so long as we provide these names to the solver when they differ from the default. The following is an example where h is used as the drift.
End of explanation
if torch.cuda.is_available():
    gpu = torch.device('cuda')
    sde = SDE().to(gpu)
    ts = ts.to(gpu)
    y0 = y0.to(gpu)
    with torch.no_grad():
        ys = torchsde.sdeint(sde, y0, ts, method='euler')  # (100, 3, 1).
    
    plot(ts, ys, xlabel='$t$', ylabel='$Y_t$')
Explanation: Trivially, the previous code may be adapted to run on GPUs, just by moving all tensors to a GPU:
End of explanation
ts = torch.linspace(0, 1, t_size)
bm = torchsde.BrownianInterval(t0=0.0, t1=1.0, size=(batch_size, state_size))
bm_increments = torch.stack([bm(t0, t1) for t0, t1 in zip(ts[:-1], ts[1:])], dim=0)
bm_queries = torch.cat((torch.zeros(1, batch_size, state_size), torch.cumsum(bm_increments, dim=0)))
plot(ts, bm_queries, xlabel='$t$', ylabel='$W_t$', title='Query')
bm_increments2 = torch.stack([bm(t0, t1) for t0, t1 in zip(ts[:-1], ts[1:])], dim=0)
bm_queries2 = torch.cat((torch.zeros(1, batch_size, state_size), torch.cumsum(bm_increments2, dim=0)))
plot(ts, bm_queries, xlabel='$t$', ylabel='$W_t$', 
     title='Query again (samples should be same as before)')
assert torch.allclose(bm_queries, bm_queries2)
Explanation: A side note is that multi-GPU data parallel is possible with the existing codebase, but the use case has not been tested out extensively and may require defining non-standard SDE classes and methods. 
2. Explicit control over randomness from the Brownian motion
To gain control over the randomness, we draw Brownian motion samples by instantiating objects of classes BrownianInterval.
End of explanation
if torch.cuda.is_available():
    bm = torchsde.BrownianInterval(t0=0.0, t1=1.0, size=(batch_size, state_size), device=gpu)
    print(bm(0.0, 0.5))
Explanation: We can also create the Brownian motion on GPUs by specifying device:
End of explanation
sde = SDE()
ts = torch.linspace(0, 1, t_size)
y0 = torch.zeros((batch_size, 1)).fill_(0.1)  # (batch_size, state_size).
bm = torchsde.BrownianInterval(t0=0.0, t1=1.0, size=(batch_size, state_size))
with torch.no_grad():
    ys = torchsde.sdeint(sde, y0, ts, method='milstein', bm=bm)
plot(ts, ys, xlabel='$t$', ylabel='$Y_t$', title='Solve SDE')
with torch.no_grad():
    ys = torchsde.sdeint(sde, y0, ts, method='milstein', bm=bm)
plot(ts, ys, xlabel='$t$', ylabel='$Y_t$', 
     title='Solve SDE again (samples should be same as before)')
# Use a new BM sample, we expect different sample paths.
bm = torchsde.BrownianInterval(t0=0.0, t1=1.0, size=(batch_size, state_size))
with torch.no_grad():
    ys = torchsde.sdeint(sde, y0, ts, method='milstein', bm=bm)
plot(ts, ys, xlabel='$t$', ylabel='$Y_t$', 
     title='Solve SDE (expect different sample paths)')
Explanation: Having a Brownian motion object helps us gain control over the randomness better. We can feed the object into the solver such that the solver's solution is conditioned on this path.
End of explanation
ys = torchsde.sdeint(sde, y0, ts, method='euler', bm=bm)
y_final = ys[-1]
target = torch.randn_like(y_final)
loss = ((target - y_final) ** 2).sum(dim=1).mean(dim=0)
loss.backward()
print(sde.theta.grad)
ys = torchsde.sdeint(sde, y0, ts, method='euler', bm=bm)
y_final = ys[-1]
target = torch.randn_like(y_final)
loss = ((target - y_final) ** 2).sum(dim=1).mean(dim=0)
grad, = torch.autograd.grad(loss, sde.theta)
print(grad)
Explanation: 3. Noise type of SDEs affects which solvers can be used and what strong orders can be attained<a id='noise_type'></a>
The supported noise types of this codebase are "diagonal", "additive", "scalar", and "general".
The following is a simple summary of each type:
- "diagonal": The diffusion function is an elementwise function, whose output has the same size as that of the state (both are of size (batch_size, state_size)). There are state_size independent Brownian motions, each responsible for the noise of only a single state dimension.
- "additive": The diffusion function is constant w.r.t. the state (but not necessarily time). The output of the diffusion function is of size (batch_size, state_size, brownian_size), and the system has brownian_size independent Brownian motions. The integral involving the Brownian motion can be loosely interpreted as integrating a sequence of matrix-vector products.
- "scalar": The diffusion function has output shape (batch_size, state_size, 1), and a single Brownian motion is shared across all state dimensions. 
- "general": The diffusion function has output shape (batch_size, state_size, brownian_size), and the system has brownian_size independent Brownian motions.
It is tempting to use the noise type configuration "general" for all problems. However, since there's little known structure for these SDEs, solvers with high strong-order cannot be constructed.
Lastly, for modeling problems, our limited experience have found "diagonal" to be a good setting, where flexibility of models and tractability of numerical integration is rather well-balanced. 
4. Computing gradients
Computing gradients through SDE solutions can be easily done by calling the backward function on loss tensors, or using torch.autograd.grad:
End of explanation
ys = torchsde.sdeint_adjoint(sde, y0, ts, method='euler', bm=bm)
y_final = ys[-1]
target = torch.randn_like(y_final)
loss = ((target - y_final) ** 2).sum(dim=1).mean(dim=0)
loss.backward()
print(sde.theta.grad)
Explanation: Switching to adjoint-mode gradient computation is as simple as replacing sdeint with sdeint_adjoint
End of explanation
<END_TASK>