jax_rb.simulation
Simulator
Simulator for global_manifold
- class jax_rb.simulation.simulator.RunParams(x_0, key, t_final, n_path, n_div, d_coeff, wiener_dim, m_size, normalize, run_type)[source]
Parameters to save a run in simulator.
- Parameters:
x_0 -- starting point of the simulation
key -- key to generate the random numbers used in simulation. Created from jax.random.PRNGKey, then jax.random.split.
t_final -- The final time of simulation. Starting time is \(t=0\).
n_path -- number of paths used in simulation
n_div -- number of subdivision (interval will be t_final/n_div
d_coeff -- difusion coefficient, d_coeff = 0.5 for the Riemannian Brownian motion.
wiener_dim -- dimension of the Wienner process used in simulation. Usually it is the dimension of the ambient space \(\mathcal{E}\). In some cases, we can simulate using the dimension of the manifold itself.
m_size -- a param indicating the size of the manifold, use to differentiate when simulating several manifolds,
normalize -- whether to normalize the move to a fixed lengh,
run_type -- string indicating one of the simulation moves. This is a tag to distinguish the output, does not affect the results.
- class jax_rb.simulation.simulator.Simulator(path_pay_off, final_pay_off)[source]
Class to do simulation on a manifold with particular funtion or simulators. Run results is saved in self.runs.
:param path_pay_off is the function value evaluated along the path :param final_pay_off is the function value evaluated at final time
- jax_rb.simulation.simulator.simulate(x_0, integrator, path_pay_off, final_pay_off, params)[source]
A simulation from \(t=0\) up to time \(t=t_final\), with time increment \(t=\frac{t_final}{n_div}\), run \(n_path\) path. Return the full distribution of the simulation. We use the minimum cut-off with accuracy level 0.5 in this version.
- Parameters:
x_0 -- starting point of the simulation
integrator -- one of the integrators (geodesic, ito, stratonovich
path_pay_off -- the cost evaluated along the path
final_pay_off -- the contribution evaluated at the final time
params -- additional parameters for the simulations: sk, t_final, n_path, n_div, d_coeff, wiener_dim
Matrix Lie Group Integrator
Module implementing simulation methods for left invariant matrix Lie group
- jax_rb.simulation.matrix_group_integrator.geodesic_move(mnf, x, unit_move, scale)[source]
\(\text{unit_move}\) is reshaped to the shape conforming with sigma., usually the shape of the ambient space. The move is \(x_{new} = \mathfrak{r}(x, \sigma(x)(\text{unit_move}(\text{scale})^{\frac{1}{2}}))\)
- jax_rb.simulation.matrix_group_integrator.geodesic_move_dim_g(mnf, x, unit_move, scale)[source]
\(\text{unit_move}\) is of dimension \(\dim \mathrm{G}\). The move is \(x_{new} = \mathfrak{r}(x, \sigma_{la}(x)(\text{unit_move}(\text{scale})^{\frac{1}{2}}))\)
- jax_rb.simulation.matrix_group_integrator.geodesic_move_dim_g_normalized(mnf, x, unit_move, scale)[source]
Similar to geodesic_move_dim_g, but unit move is rescaled to have fixed length 1 in the metric of the group.
- jax_rb.simulation.matrix_group_integrator.geodesic_move_normalized(mnf, x, unit_move, scale)[source]
Similar to geodesic_move, but unit move is rescaled to have fixed length 1 in the metric of the group.
- jax_rb.simulation.matrix_group_integrator.geodesic_move_with_drift(mnf, x, unit_move, scale, id_additional_drift)[source]
This method is used to simulate a Riemanian Brownian motion with drift. \(\text{unit_move}\) is reshaped to the shape conforming with sigma., usually the shape of the ambient space. \(\text{id_additional_drift}\) is an element of the Lie algebra. The move is \(x_{new} = \mathfrak{r}(x, \sigma(x)((\text{scale})^{\frac{1}{2}}\times \text{unit_move})+\text{scale}\times x (\text{id_additional_drift}))\)
- jax_rb.simulation.matrix_group_integrator.ito_move_dim_g(mnf, x, unit_move, scale)[source]
Similar to rbrownian_ito_move, but driven with a Wiener process of dimension \(\dim \mathrm{G}\).
- jax_rb.simulation.matrix_group_integrator.ito_move_with_drift(mnf, x, unit_move, scale, id_additional_drift)[source]
This method is used to simulate a Riemanian Brownian motion with drift given in Ito form. Use stochastic projection method to solve the Ito equation. The drift is given as an element of the Lie algebra. Use Euler Maruyama here.
- jax_rb.simulation.matrix_group_integrator.rbrownian_ito_move(mnf, x, unit_move, scale)[source]
Use stochastic projection method to solve the Ito equation. Use Euler Maruyama here.
- jax_rb.simulation.matrix_group_integrator.rbrownian_stratonovich_move(mnf, x, unit_move, scale)[source]
Using projection method to solve the Stratonovich equation. In many cases \(v_0\) is zero (unimodular group). Use Euler Heun.
- jax_rb.simulation.matrix_group_integrator.stratonovich_move_dim_g(mnf, x, unit_move, scale)[source]
Similar to rbrownian_stratonovich_move, but driven with a Wiener process of dimension \(\dim \mathrm{G}\).
- jax_rb.simulation.matrix_group_integrator.stratonovich_move_with_drift(mnf, x, unit_move, scale, id_additional_drift)[source]
This method is used to simulate a Riemanian Brownian motion with drift given in Stratonovich form. Using projection method to solve the Stratonovich equation. The additional drift is on top of the RB term, given as an element of the Lie algebra In many cases \(v_0\) is zero (unimodular group). Use Euler Heun.
Global Manifold Integrator
Module implementing simulation methods for embedded manifolds
- jax_rb.simulation.global_manifold_integrator.geodesic_move(mnf, x, unit_move, scale)[source]
simulate using a second order retraction. The move is \(x_{new} = \mathfrak{r}(x, \Pi(x)\sigma(x)(\text{unit_move}(\text{scale})^{\frac{1}{2}}))\)
- jax_rb.simulation.global_manifold_integrator.geodesic_move_exact(mnf, x, unit_move, scale)[source]
similar to geodesic_move, but use exact geodesic
- jax_rb.simulation.global_manifold_integrator.geodesic_move_exact_normalized(mnf, x, unit_move, scale)[source]
similar to geodesic_move_exact, but use normalize the unit_move
- jax_rb.simulation.global_manifold_integrator.geodesic_move_normalized(mnf, x, unit_move, scale)[source]
similar to geodesic_move, but the move is normalized to have fixed length \(scale^{\frac{1}{2}}\)
- jax_rb.simulation.global_manifold_integrator.geodesic_move_with_drift(mnf, x, unit_move, scale, additional_drift)[source]
This method is used to simulate a Riemanian Brownian motion with drift. The additional_drift is on top of the Brownian motion. Simulate using a second order retraction. The move is \(x_{new} = \mathfrak{r}(x, \Pi(x)\sigma(x)(\text{unit_move}(\text{scale})^{\frac{1}{2}}+\text{scale}\times\text{additional_drift}))\)
- jax_rb.simulation.global_manifold_integrator.ito_move_with_drift(mnf, x, unit_move, scale, additional_drift)[source]
This method is used to simulate a Riemanian Brownian motion with drift. The additional_drift is on top of the Brownian motion. Use Euler Maruyama and projection method to solve the Ito equation.
- jax_rb.simulation.global_manifold_integrator.rbrownian_ito_move(mnf, x, unit_move, scale)[source]
Use Euler Maruyama and projection method to solve the Ito equation.
- jax_rb.simulation.global_manifold_integrator.rbrownian_stratonovich_move(mnf, x, unit_move, scale)[source]
Use Euler Heun and projection method to solve the Stratonovich equation.
- jax_rb.simulation.global_manifold_integrator.stratonovich_move_with_drift(mnf, x, unit_move, scale, additional_drift)[source]
This method is used to simulate a Riemanian Brownian motion with drift. The additional_drift is on top of the Brownian motion. Use Euler Heun and projection method to solve the Stratonovich equation.
Retractive Integrator
Module implementing the retractive Euler-Maruyama integrator.
- jax_rb.simulation.retractive_integrator.retractive_move(rtr, x, t, unit_move, scale, sigma, mu)[source]
Simulating the equation \(dX_t = \mu(X_t, t) dt + \sigma(X_t, t) dW_t\) using the retraction rtr. We do not assume a Riemanian metric on the manifold, \(\sigma\sigma^T\) could be degenerated on \(T\mathcal{M}\).
W is a Wiener process driving the equation, defined on \(\mathbb{R}^k\). W is given by unit_move.
\(\sigma(X_t, t)\) maps \(\mathbb{R}^k\) to \(\mathcal{E}\), but the image belongs to \(T_{X_t}\mathcal{M}\).
The retraction rtr is assume to have the method \(\text{drift_adj}\) for an adjustment.
The move is \(x_{new} = \mathfrak{r}(x, \Pi(x)\sigma(x)(\text{unit_move}(\text{scale})^{\frac{1}{2}}) + \text{scale} (\mu + \text{drift_adj}))\).
- Parameters:
rtr -- the retraction,
x -- a point on the manifold,
t -- time
unit_move -- a random normal draw
scale -- scaling
sigma -- a function implementing the map \(\sigma\)
mu -- a function implementing the Ito drift \(\mu\)