Repository: pfnet-research/pfhedge (language: python)

Issue:
Not accurate calculation of entripic_risk_measure when the value is significantlly big
```
from pfhedge.nn import EntropicRiskMeasure

loss = EntropicRiskMeasure()
input = -torch.arange(4.0)-100
loss(input)
＃tensor(inf)
```
ref.: https://pfnet-research.github.io/pfhedge/generated/pfhedge.nn.EntropicRiskMeasure.html

This result has to be between -100 and -104.
This calculation accuracy can be improved in the same manner as the usual softmax calculation.



Discussion / hints:

Fix patch:
diff --git a/docs/source/instruments.rst b/docs/source/instruments.rst
index fb0cf77f..35e6686b 100644
--- a/docs/source/instruments.rst
+++ b/docs/source/instruments.rst
@@ -31,6 +31,7 @@ Primary Instruments
     instruments.BrownianStock
     instruments.HestonStock
     instruments.LocalVolatilityStock
+    instruments.MertonJumpStock
     instruments.RoughBergomiStock
     instruments.CIRRate
     instruments.VasicekRate
diff --git a/docs/source/stochastic.rst b/docs/source/stochastic.rst
index e8bda60c..2974fcd8 100644
--- a/docs/source/stochastic.rst
+++ b/docs/source/stochastic.rst
@@ -19,6 +19,7 @@ Stochastic Processes
     generate_heston
     generate_local_volatility_process
     generate_vasicek
+    generate_merton_jump
     generate_rough_bergomi
 
 Generators
diff --git a/pfhedge/instruments/__init__.py b/pfhedge/instruments/__init__.py
index 4b249cde..226c927f 100644
--- a/pfhedge/instruments/__init__.py
+++ b/pfhedge/instruments/__init__.py
@@ -16,5 +16,6 @@
 from .primary.cir import CIRRate  # NOQA
 from .primary.heston import HestonStock  # NOQA
 from .primary.local_volatility import LocalVolatilityStock  # NOQA
+from .primary.merton_jump import MertonJumpStock  # NOQA
 from .primary.rough_bergomi import RoughBergomiStock  # NOQA
 from .primary.vasicek import VasicekRate  # NOQA
diff --git a/pfhedge/instruments/primary/merton_jump.py b/pfhedge/instruments/primary/merton_jump.py
new file mode 100644
index 00000000..3384c2ce
--- /dev/null
+++ b/pfhedge/instruments/primary/merton_jump.py
@@ -0,0 +1,184 @@
+from math import ceil
+from typing import Callable
+from typing import Optional
+from typing import Tuple
+
+import torch
+from torch import Tensor
+
+from pfhedge._utils.doc import _set_attr_and_docstring
+from pfhedge._utils.doc import _set_docstring
+from pfhedge._utils.str import _format_float
+from pfhedge._utils.typing import TensorOrScalar
+from pfhedge.stochastic import generate_merton_jump
+
+from .base import BasePrimary
+
+
+class MertonJumpStock(BasePrimary):
+    r"""A stock of which spot price and variance follow Merton Jump Diffusion process.
+
+    .. seealso::
+        - :func:`pfhedge.stochastic.generate_merton_jump`:
+          The stochastic process.
+
+    Args:
+        mu (float, default=0.0): The drift parameter :math:`\mu`.
+        sigma (float, default=0.2): The volatility parameter :math:`\sigma`.
+        jump_per_year (float, default=1.0): The average number of jumps per year.
+        jump_mean (float, default=0.0): The mean of jumnp sizes.
+        jump_std (float, default=0.3): The standard deviation of jump sizes.
+        cost (float, default=0.0): The transaction cost rate.
+        dt (float, default=1/250): The time step interval.
+        dtype (torch.device, optional): Desired device of returned tensor.
+            Default: If None, uses a global default
+            (see :func:`torch.set_default_tensor_type()`).
+        device (torch.device, optional): Desired device of returned tensor.
+            Default: if None, uses the current device for the default tensor type
+            (see :func:`torch.set_default_tensor_type()`).
+            ``device`` will be the CPU for CPU tensor types and
+            the current CUDA device for CUDA tensor types.
+        engine (callable, default=torch.randn): The desired generator of random numbers
+            from a standard normal distribution.
+            A function call ``engine(size, dtype=None, device=None)``
+            should return a tensor filled with random numbers
+            from a standard normal distribution.
+
+    Buffers:
+        - spot (:class:`torch.Tensor`): The spot price of the instrument.
+          This attribute is set by the method :meth:`simulate()`.
+          The shape is :math:`(N, T)` where
+          :math:`N` is the number of simulated paths and
+          :math:`T` is the number of time steps.
+        - variance (:class:`torch.Tensor`): The variance of the instrument.
+          Note that this is different from the realized variance of the spot price.
+          This attribute is set by the method :meth:`simulate()`.
+          The shape is :math:`(N, T)`.
+
+    Examples:
+        >>> from pfhedge.instruments import MertonJumpStock
+        >>>
+        >>> _ = torch.manual_seed(42)
+        >>> stock = MertonJumpStock()
+        >>> stock.simulate(n_paths=2, time_horizon=5/250)
+        >>> stock.spot
+        tensor([[1.0000, 1.0100, 1.0135, 1.0141, 1.0208, 1.0176],
+                [1.0000, 1.0066, 0.9955, 1.0047, 1.0063, 1.0172]])
+        >>> stock.variance
+        tensor([[0.0400, 0.0400, 0.0400, 0.0400, 0.0400, 0.0400],
+                [0.0400, 0.0400, 0.0400, 0.0400, 0.0400, 0.0400]])
+        >>> stock.volatility
+        tensor([[0.2000, 0.2000, 0.2000, 0.2000, 0.2000, 0.2000],
+                [0.2000, 0.2000, 0.2000, 0.2000, 0.2000, 0.2000]])
+    """
+
+    spot: Tensor
+
+    def __init__(
+        self,
+        mu: float = 0.0,
+        sigma: float = 0.2,
+        jump_per_year: float = 68,
+        jump_mean: float = 0.0,
+        jump_std: float = 0.01,
+        cost: float = 0.0,
+        dt: float = 1 / 250,
+        dtype: Optional[torch.dtype] = None,
+        device: Optional[torch.device] = None,
+        engine: Callable[..., Tensor] = torch.randn,
+    ) -> None:
+        super().__init__()
+
+        self.mu = mu
+        self.sigma = sigma
+        self.jump_per_year = jump_per_year
+        self.jump_mean = jump_mean
+        self.jump_std = jump_std
+        self.cost = cost
+        self.dt = dt
+        self.engine = engine
+
+        self.to(dtype=dtype, device=device)
+
+    @property
+    def default_init_state(self) -> Tuple[float, ...]:
+        return (1.0,)
+
+    @property
+    def volatility(self) -> Tensor:
+        """Returns the volatility of self.
+
+        It is a tensor filled with ``self.sigma``.
+        """
+        return torch.full_like(self.get_buffer("spot"), self.sigma)
+
+    @property
+    def variance(self) -> Tensor:
+        """Returns the volatility of self.
+
+        It is a tensor filled with the square of ``self.sigma``.
+        """
+        return torch.full_like(self.get_buffer("spot"), self.sigma ** 2)
+
+    def simulate(
+        self,
+        n_paths: int = 1,
+        time_horizon: float = 20 / 250,
+        init_state: Optional[Tuple[TensorOrScalar, ...]] = None,
+    ) -> None:
+        """Simulate the spot price and add it as a buffer named ``spot``.
+
+        The shape of the spot is :math:`(N, T)`, where
+        :math:`N` is the number of simulated paths and
+        :math:`T` is the number of time steps.
+        The number of time steps is determinded from ``dt`` and ``time_horizon``.
+
+        Args:
+            n_paths (int, default=1): The number of paths to simulate.
+            time_horizon (float, default=20/250): The period of time to simulate
+                the price.
+            init_state (tuple[torch.Tensor | float], optional): The initial
+                state of the instrument.
+                This is specified by a tuple :math:`(S(0), V(0))` where
+                :math:`S(0)` and :math:`V(0)` are the initial values of
+                spot and variance, respectively.
+                If ``None`` (default), it uses the default value
+                (See :attr:`default_init_state`).
+        """
+        if init_state is None:
+            init_state = self.default_init_state
+
+        output = generate_merton_jump(
+            n_paths=n_paths,
+            n_steps=ceil(time_horizon / self.dt + 1),
+            init_state=init_state,
+            sigma=self.sigma,
+            mu=self.mu,
+            jump_per_year=self.jump_per_year,
+            jump_mean=self.jump_mean,
+            jump_std=self.jump_std,
+            dt=self.dt,
+            dtype=self.dtype,
+            device=self.device,
+            engine=self.engine,
+        )
+
+        self.register_buffer("spot", output)
+
+    def extra_repr(self) -> str:
+        params = [
+            "mu=" + _format_float(self.mu),
+            "sigma=" + _format_float(self.sigma),
+            "jump_per_year=" + _format_float(self.jump_per_year),
+            "jump_mean=" + _format_float(self.jump_mean),
+            "jump_std=" + _format_float(self.jump_std),
+        ]
+        if self.cost != 0.0:
+            params.append("cost=" + _format_float(self.cost))
+        params.append("dt=" + _format_float(self.dt))
+        return ", ".join(params)
+
+
+# Assign docstrings so they appear in Sphinx documentation
+_set_docstring(MertonJumpStock, "default_init_state", BasePrimary.default_init_state)
+_set_attr_and_docstring(MertonJumpStock, "to", BasePrimary.to)
diff --git a/pfhedge/stochastic/__init__.py b/pfhedge/stochastic/__init__.py
index 82cc15de..3861aae0 100644
--- a/pfhedge/stochastic/__init__.py
+++ b/pfhedge/stochastic/__init__.py
@@ -3,6 +3,7 @@
 from .cir import generate_cir  # NOQA
 from .heston import generate_heston  # NOQA
 from .local_volatility import generate_local_volatility_process  # NOQA
+from .merton_jump import generate_merton_jump  # NOQA
 from .random import randn_antithetic  # NOQA
 from .random import randn_sobol_boxmuller  # NOQA
 from .rough_bergomi import generate_rough_bergomi  # NOQA
diff --git a/pfhedge/stochastic/merton_jump.py b/pfhedge/stochastic/merton_jump.py
new file mode 100644
index 00000000..5e62066e
--- /dev/null
+++ b/pfhedge/stochastic/merton_jump.py
@@ -0,0 +1,113 @@
+import math
+from typing import Callable
+from typing import Optional
+from typing import Tuple
+from typing import Union
+
+import torch
+from torch import Tensor
+
+from pfhedge._utils.typing import TensorOrScalar
+from pfhedge.stochastic._utils import cast_state
+
+
+def generate_merton_jump(
+    n_paths: int,
+    n_steps: int,
+    init_state: Union[Tuple[TensorOrScalar, ...], TensorOrScalar] = (1.0,),
+    mu: float = 0.0,
+    sigma: float = 0.2,
+    jump_per_year: float = 68.2,
+    jump_mean: float = 0.0,
+    jump_std: float = 0.02,
+    dt: float = 1 / 250,
+    dtype: Optional[torch.dtype] = None,
+    device: Optional[torch.device] = None,
+    engine: Callable[..., Tensor] = torch.randn,
+) -> Tensor:
+    r"""Returns time series following the Merton's Jump Diffusion Model.
+
+    The time evolution of the process is given by:
+
+    .. math::
+
+        \frac{dS(t)}{S(t)} = (\mu - \lambda k) dt + \sigma dW(t) + dJ(t) .
+
+    Reference:
+     - Merton, R.C., Option pricing when underlying stock returns are discontinuous,
+       Journal of Financial Economics, 3 (1976), 125-144.
+     - Gugole, N. (2016). Merton jump-diffusion model versus the black and scholes approach for the log-returns and volatility
+       smile fitting. International Journal of Pure and Applied Mathematics, 109(3), 719-736.
+       (Especially p.730 for parameter calibulation.)
+
+    Args:
+        n_paths (int): The number of simulated paths.
+        n_steps (int): The number of time steps.
+        init_state (tuple[torch.Tensor | float], default=(1.0,)): The initial state of
+            the time series.
+            This is specified by a tuple :math:`(S(0),)`.
+            It also accepts a :class:`torch.Tensor` or a :class:`float`.
+        mu (float, default=0.0): The parameter :math:`\mu`,
+            which stands for the dirft coefficient of the time series.
+        sigma (float, default=0.2): The parameter :math:`\sigma`,
+            which stands for the volatility of the time series.
+        jump_per_year (float, default=68.2): The average number of jumps per year.
+        jump_mean (float, default=0.0): The mean of jump sizes.
+        jump_std (float, default=0.02): The standard deviation of jump sizes.
+        dt (float, default=1/250): The time step interval.
+        dtype (torch.dtype, optional): The desired data type of returned tensor.
+            Default: If ``None``, uses a global default
+            (see :func:`torch.set_default_tensor_type()`).
+        device (torch.device, optional): The desired device of returned tensor.
+            Default: If ``None``, uses the current device for the default tensor type
+            (see :func:`torch.set_default_tensor_type()`).
+            ``device`` will be the CPU for CPU tensor types and the current CUDA device
+            for CUDA tensor types.
+        engine (callable, default=torch.randn): The desired generator of random numbers
+            from a standard normal distribution.
+            A function call ``engine(size, dtype=None, device=None)``
+            should return a tensor filled with random numbers
+            from a standard normal distribution.
+
+    Shape:
+        - Output: :math:`(N, T)` where
+          :math:`N` is the number of paths and
+          :math:`T` is the number of time steps.
+
+    Returns:
+        torch.Tensor
+
+    Examples:
+        >>> from pfhedge.stochastic import generate_merton_jump
+        >>>
+        >>> _ = torch.manual_seed(42)
+        >>> generate_merton_jump(2, 5)
+        tensor([[1.0000, 0.9904, 1.0074, 1.0160, 1.0117],
+                [1.0000, 1.0034, 1.0040, 1.0490, 1.0457]])
+    """
+    # https://www.codearmo.com/python-tutorial/merton-jump-diffusion-model-python
+    init_state = cast_state(init_state, dtype=dtype, device=device)
+
+    poisson = torch.distributions.poisson.Poisson(rate=jump_per_year * dt)
+    n_jumps = poisson.sample((n_paths, n_steps - 1)).to(dtype=dtype, device=device)
+
+    # Eq. (3) in https://www.imes.boj.or.jp/research/papers/japanese/kk22-b1-3.pdf
+    jump = (
+        jump_mean * n_jumps
+        + engine(*(n_paths, n_steps - 1), dtype=dtype, device=device)
+        * jump_std
+        * n_jumps.sqrt()
+    )
+    jump = torch.cat(
+        [torch.zeros((n_paths, 1), dtype=dtype, device=device), jump], dim=1
+    )
+
+    randn = engine(*(n_paths, n_steps), dtype=dtype, device=device)
+    randn[:, 0] = 0.0
+    drift = (
+        (mu - (sigma ** 2) / 2 - jump_per_year * (jump_mean + jump_std ** 2 / 2))
+        * dt
+        * torch.arange(n_steps).to(randn)
+    )
+    brown = randn.cumsum(1) * math.sqrt(dt)
+    return init_state[0] * (drift + sigma * brown + jump.cumsum(1)).exp()


Test patch:
diff --git a/tests/instruments/primary/test_merton_jump.py b/tests/instruments/primary/test_merton_jump.py
new file mode 100644
index 00000000..d205c2d7
--- /dev/null
+++ b/tests/instruments/primary/test_merton_jump.py
@@ -0,0 +1,41 @@
+import pytest
+import torch
+
+from pfhedge.instruments import MertonJumpStock
+
+
+class TestMertonJumpStock:
+    @pytest.mark.parametrize("seed", range(1))
+    def test_values_are_finite(self, seed, device: str = "cpu"):
+        torch.manual_seed(seed)
+
+        s = MertonJumpStock().to(device)
+        s.simulate(n_paths=1000)
+
+        assert not s.variance.isnan().any()
+
+    @pytest.mark.gpu
+    @pytest.mark.parametrize("seed", range(1))
+    def test_values_are_finite_gpu(self, seed):
+        self.test_values_are_finite(seed, device="cuda")
+
+    def test_repr(self):
+        s = MertonJumpStock(cost=1e-4)
+        expect = "MertonJumpStock(\
+mu=0., sigma=0.2000, jump_per_year=68, jump_mean=0., jump_std=0.0100, cost=1.0000e-04, dt=0.0040)"
+        assert repr(s) == expect
+
+    def test_simulate_shape(self, device: str = "cpu"):
+        s = MertonJumpStock(dt=0.1).to(device)
+        s.simulate(time_horizon=0.2, n_paths=10)
+        assert s.spot.size() == torch.Size((10, 3))
+        assert s.variance.size() == torch.Size((10, 3))
+
+        s = MertonJumpStock(dt=0.1).to(device)
+        s.simulate(time_horizon=0.25, n_paths=10)
+        assert s.spot.size() == torch.Size((10, 4))
+        assert s.variance.size() == torch.Size((10, 4))
+
+    @pytest.mark.gpu
+    def test_simulate_shape_gpu(self):
+        self.test_simulate_shape(device="cuda")
diff --git a/tests/stochastic/test_merton_jump.py b/tests/stochastic/test_merton_jump.py
new file mode 100644
index 00000000..88b06beb
--- /dev/null
+++ b/tests/stochastic/test_merton_jump.py
@@ -0,0 +1,283 @@
+from math import sqrt
+
+import pytest
+import torch
+from torch.testing import assert_close
+
+from pfhedge.stochastic import generate_merton_jump
+from pfhedge.stochastic.engine import RandnSobolBoxMuller
+
+
+def test_generate_brownian_mean_no_jump(device: str = "cpu"):
+    torch.manual_seed(42)
+    n_paths = 10000
+    n_steps = 250
+
+    output = generate_merton_jump(
+        n_paths, n_steps, jump_std=0.0, device=torch.device(device)
+    )
+    assert output.size() == torch.Size((n_paths, n_steps))
+    result = output[:, -1].mean()
+    expect = torch.ones_like(result)
+    std = 0.2 * sqrt(1 / n_paths)
+    assert_close(result, expect, atol=3 * std, rtol=0)
+
+
+@pytest.mark.gpu
+def test_generate_brownian_mean_no_jump_gpu():
+    test_generate_brownian_mean_no_jump(device="cuda")
+
+
+def test_generate_brownian_mean_no_jump1(device: str = "cpu"):
+    torch.manual_seed(42)
+    n_paths = 10000
+    n_steps = 250
+
+    output = generate_merton_jump(n_paths, n_steps, jump_per_year=0.0, device=device)
+    assert output.size() == torch.Size((n_paths, n_steps))
+    result = output[:, -1].mean()
+    expect = torch.ones_like(result)
+    std = 0.2 * sqrt(1 / n_paths)
+    assert_close(result, expect, atol=3 * std, rtol=0)
+
+
+@pytest.mark.gpu
+def test_generate_brownian_mean_no_jump1_gpu():
+    test_generate_brownian_mean_no_jump1(device="cuda")
+
+
+def test_generate_brownian_mean_no_jump_std(device: str = "cpu"):
+    torch.manual_seed(42)
+    n_paths = 10000
+    n_steps = 250
+
+    output = generate_merton_jump(
+        n_paths,
+        n_steps,
+        jump_per_year=1000,
+        jump_std=0.0,
+        jump_mean=0.1,
+        device=torch.device(device),
+    )
+    assert output.size() == torch.Size((n_paths, n_steps))
+    result = output[:, -1].log().mean()
+    expect = torch.zeros_like(result)
+    std = 0.2 * sqrt(1 / n_paths)
+    assert_close(result, expect, atol=3 * std, rtol=0)
+
+
+@pytest.mark.gpu
+def test_generate_brownian_mean_no_jump_std_gpu():
+    test_generate_brownian_mean_no_jump_std(device="cuda")
+
+
+def test_generate_brownian_mean(device: str = "cpu"):
+    torch.manual_seed(42)
+    n_paths = 10000
+    n_steps = 250
+
+    output = generate_merton_jump(
+        n_paths, n_steps, jump_per_year=1, device=torch.device(device)
+    )
+    assert output.size() == torch.Size((n_paths, n_steps))
+    result = output[:, -1].mean()
+    expect = torch.ones_like(result)
+    std = 0.2 * sqrt(1 / n_paths) + 0.3 * sqrt(1 / n_paths)
+    assert_close(result, expect, atol=3 * std, rtol=0)
+
+
+@pytest.mark.gpu
+def test_generate_brownian_mean_gpu():
+    test_generate_brownian_mean(device="cuda")
+
+
+def test_generate_merton_jump_nosigma(device: str = "cpu"):
+    torch.manual_seed(42)
+    n_steps = 250
+
+    result = generate_merton_jump(
+        1, n_steps, sigma=0, jump_per_year=0, device=torch.device(device)
+    )
+    expect = torch.ones(1, n_steps).to(device)
+    assert_close(result, expect)
+
+    mu = 0.1
+    dt = 0.01
+    result = generate_merton_jump(
+        1, n_steps, mu=mu, sigma=0, dt=dt, jump_per_year=0, device=torch.device(device)
+    ).log()
+    expect = torch.linspace(0, mu * dt * (n_steps - 1), n_steps).unsqueeze(0).to(device)
+    assert_close(result, expect)
+
+
+@pytest.mark.gpu
+def test_generate_merton_jump_nosigma_gpu():
+    test_generate_merton_jump_nosigma(device="cpu")
+
+
+def test_generate_merton_jump_nosigma2(device: str = "cpu"):
+    torch.manual_seed(42)
+    n_steps = 250
+
+    result = generate_merton_jump(
+        1, n_steps, sigma=0, jump_std=0, device=torch.device(device)
+    )
+    expect = torch.ones(1, n_steps).to(device)
+    assert_close(result, expect)
+
+    mu = 0.1
+    dt = 0.01
+    result = generate_merton_jump(
+        1, n_steps, mu=mu, sigma=0, dt=dt, jump_std=0, device=torch.device(device)
+    ).log()
+    expect = torch.linspace(0, mu * dt * (n_steps - 1), n_steps).unsqueeze(0).to(device)
+    assert_close(result, expect)
+
+
+@pytest.mark.gpu
+def test_generate_merton_jump_nosigma2_gpu():
+    test_generate_merton_jump_nosigma2(device="cuda")
+
+
+def test_generate_merton_jump_std(device: str = "cpu"):
+    torch.manual_seed(42)
+    n_paths = 10000
+    n_steps = 250
+
+    output = generate_merton_jump(
+        n_paths, n_steps, jump_per_year=0, device=torch.device(device)
+    )
+    assert output.size() == torch.Size((n_paths, n_steps))
+    result = output[:, -1].log().std()
+    expect = torch.full_like(result, 0.2)
+    assert_close(result, expect, atol=0, rtol=0.1)
+
+
+@pytest.mark.gpu
+def test_generate_merton_jump_std_gpu():
+    test_generate_merton_jump_std(device="cuda")
+
+
+def test_generate_merton_jump_std2(device: str = "cpu"):
+    torch.manual_seed(42)
+    n_paths = 10000
+    n_steps = 250
+
+    output = generate_merton_jump(
+        n_paths, n_steps, jump_std=0, device=torch.device(device)
+    )
+    assert output.size() == torch.Size((n_paths, n_steps))
+    result = output[:, -1].log().std()
+    expect = torch.full_like(result, 0.2)
+    assert_close(result, expect, atol=0, rtol=0.1)
+
+
+@pytest.mark.gpu
+def test_generate_merton_jump_std2_gpu():
+    test_generate_merton_jump_std2(device="cuda")
+
+
+def test_generate_merton_jump_mean_init_state(device: str = "cpu"):
+    torch.manual_seed(42)
+    n_paths = 10000
+    n_steps = 250
+
+    output = generate_merton_jump(
+        n_paths, n_steps, init_state=1.0, jump_per_year=0, device=torch.device(device)
+    )
+    assert output.size() == torch.Size((n_paths, n_steps))
+    result = output[:, -1].mean()
+    expect = torch.ones_like(result)
+    std = 0.2 * sqrt(1 / n_paths)
+    assert_close(result, expect, atol=3 * std, rtol=0)
+
+    output = generate_merton_jump(
+        n_paths,
+        n_steps,
+        init_state=torch.tensor(1.0),
+        jump_per_year=0,
+        device=torch.device(device),
+    )
+    assert output.size() == torch.Size((n_paths, n_steps))
+    result = output[:, -1].mean()
+    expect = torch.ones_like(result)
+    std = 0.2 * sqrt(1 / n_paths)
+    assert_close(result, expect, atol=3 * std, rtol=0)
+
+    output = generate_merton_jump(
+        n_paths,
+        n_steps,
+        init_state=torch.tensor([1.0]),
+        jump_per_year=0,
+        device=torch.device(device),
+    )
+    assert output.size() == torch.Size((n_paths, n_steps))
+    result = output[:, -1].mean()
+    expect = torch.ones_like(result)
+    std = 0.2 * sqrt(1 / n_paths)
+    assert_close(result, expect, atol=3 * std, rtol=0)
+
+
+@pytest.mark.gpu
+def test_generate_merton_jump_mean_init_state_gpu():
+    test_generate_merton_jump_mean_init_state(device="cuda")
+
+
+def test_generate_merton_jump_mean_mu(device: str = "cpu"):
+    torch.manual_seed(42)
+    n_paths = 10000
+    n_steps = 250
+    dt = 1 / 250
+    mu = 0.1
+
+    output = generate_merton_jump(
+        n_paths, n_steps, mu=mu, jump_per_year=0, device=torch.device(device)
+    )
+    result = output[:, -1].mean().log()
+    expect = torch.full_like(result, mu * dt * n_steps).to(device)
+    std = 0.2 * sqrt(1 / n_paths)
+    assert_close(result, expect, atol=3 * std, rtol=0)
+
+
+@pytest.mark.gpu
+def test_generate_merton_jump_mean_mu_gpu():
+    test_generate_merton_jump_mean_mu(device="cuda")
+
+
+def test_generate_merton_jump_dtype(device: str = "cpu"):
+    torch.manual_seed(42)
+
+    output = generate_merton_jump(
+        1, 1, dtype=torch.float32, device=torch.device(device)
+    )
+    assert output.dtype == torch.float32
+
+    output = generate_merton_jump(
+        1, 1, dtype=torch.float64, device=torch.device(device)
+    )
+    assert output.dtype == torch.float64
+
+
+@pytest.mark.gpu
+def test_generate_merton_jump_dtype_gpu():
+    test_generate_merton_jump_dtype(device="cuda")
+
+
+def test_generate_merton_jump_sobol_mean(device: str = "cpu"):
+    n_paths = 10000
+    n_steps = 250
+
+    engine = RandnSobolBoxMuller(seed=42, scramble=True)
+    output = generate_merton_jump(
+        n_paths, n_steps, engine=engine, jump_per_year=0, device=torch.device(device)
+    )
+    assert output.size() == torch.Size((n_paths, n_steps))
+    result = output[:, -1].log().mean()
+    expect = torch.zeros_like(result).to(device)
+    std = 0.2 * sqrt(1 / n_paths)
+    assert_close(result, expect, atol=10 * std, rtol=0)
+
+
+@pytest.mark.gpu
+def test_generate_merton_jump_sobol_mean_gpu():
+    test_generate_merton_jump_sobol_mean(device="cuda")