<commit_before>"""
.. module:: integrate.py
   :synopsis: Integrate stochastic master equations in vectorized form.
.. moduleauthor:: <NAME> <<EMAIL>>

"""

import numpy as np

def milstein(drift, diffusion, X0, ts, dws):
    """Integrate a system of ordinary stochastic differential equations:

    .. math::

        d\vec{X}=a(X,t)\,dt+b(X,t)\,dW_t

    Uses the Milstein method:

    .. math::

        w_{i+1}=w_i+a(w_i,t_i)\Delta t_i+b(w_i,t_i)\Delta W_i+
        \frac{1}{2}b(w_i,t_i)\frac{\partial b}{\partial X}(w_i,t_i)
        ((\Delta W_i)^2-\Delta t_i)

    :param drift:       Computes the drift term :math:`a(t,X)` at t0
    :type drift:        callable(X, t0)
    :param diffusion:   Computes the diffusion term :math:`b(t,X)` at t0
    :type diffusion:    callable(X, t0)
    :param X0:          Initial condition on X (can be a vector)
    :type X0:           array
    :param ts:          A sequence of time points for which to solve for y. The
                        initial value point should be the first element of this
                        sequence.
    :type ts:           array
    :param dws:         Normalized Weiner increments for each time step (i.e.
                        samples from a Gaussian distribution with mean 0 and
                        variance 1).
    :type dws:          array, shape (len(t) - 1)
    :return:            Array containint the value of X for each desired time in
                        t, with the initial value `X0` in the first row.
    :rtype:             numpy.array, shape (len(t), len(X0))

    """

    dts = [tf - ti for tf, ti in zip(t[1:], t[:-1])]
    # Scale the Weiner increments to the time increments.
    sqrtdt = np.sqrt(dt)
    dWs = np.product(np.array([sqrtdt, dws]), axis=0)

    X = [np.array(X0)]

    for t, dt, dw, sqrtdt in zip(ts[1:], dts):
        X.append(X[-1] + drift(X[-1], t)*dt + diffusion(X[-1], t)*sqrtdt*dw)<commit_msg>Finished writing the Miltein expression.<commit_after>"""
.. module:: integrate.py
   :synopsis: Integrate stochastic master equations in vectorized form.
.. moduleauthor:: <NAME> <<EMAIL>>

"""

import numpy as np

def milstein(drift, diffusion, diffusion_prime, X0, ts, dws):
    """Integrate a system of ordinary stochastic differential equations:

    .. math::

        d\vec{X}=a(X,t)\,dt+b(X,t)\,dW_t

    Uses the Milstein method:

    .. math::

        w_{i+1}=w_i+a(w_i,t_i)\Delta t_i+b(w_i,t_i)\Delta W_i+
        \frac{1}{2}b(w_i,t_i)\frac{\partial b}{\partial X}(w_i,t_i)
        ((\Delta W_i)^2-\Delta t_i)

    :param drift:           Computes the drift term :math:`a(t,X)` at t0
    :type drift:            callable(X, t0)
    :param diffusion:       Computes the diffusion term :math:`b(t,X)` at t0
    :type diffusion:        callable(X, t0)
    :param diffusion_prime: Computes the diffusion term :math:`b(t,X)` at t0
    :type diffusion_prime:  callable(X, t0)
    :param X0:              Initial condition on X (can be a vector)
    :type X0:               array
    :param ts:              A sequence of time points for which to solve for y.
                            The
                            initial value point should be the first element of
                            this sequence.
    :type ts:               array
    :param dws:             Normalized Weiner increments for each time step
                            (i.e. samples from a Gaussian distribution with mean
                            0 and variance 1).
    :type dws:              array, shape (len(t) - 1)
    :return:                Array containint the value of X for each desired
                            time in t, with the initial value `X0` in the first
                            row.
    :rtype:                 numpy.array, shape (len(t), len(X0))

    """

    dts = [tf - ti for tf, ti in zip(t[1:], t[:-1])]
    # Scale the Weiner increments to the time increments.
    sqrtdts = np.sqrt(dts)
    dWs = np.product(np.array([sqrtdt, dws]), axis=0)

    X = [np.array(X0)]

    for t, dt, dW in zip(ts[:-1], dts, dWs):
        X.append(X[-1] + drift(X[-1], t)*dt + diffusion(X[-1], t)*(dW +
            0.5*diffusion_prime(X[-1], t)*(dW**2 - dt)))

    return X
{'fraction_non_alphanumeric': 0.1267538644470868, 'fraction_numerical': 0.012841854934601665, 'mean_word_length': 1.9351011863224006, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4830423', 'n_tokens_mistral': 1425, 'n_tokens_neox': 1362, 'n_words': 436}