From 4156d4d669c10a5f56efc63048a4ef9033bcb770 Mon Sep 17 00:00:00 2001
From: rhandal-pfn <168500787+rhandal-pfn@users.noreply.github.com>
Date: Fri, 30 Aug 2024 15:03:59 +0900
Subject: [PATCH] kou's jump model (#633)

* kou's jump model

* test and lit fix

* tet and lint fix 2

* test and isort  fix

* distribution generator fix

* example fix

* fix engine and example on cpu

* exaple fix for old torch

* example jump change

* change annual jumps in generate example

* jump mu instead of eta

* rebase and lint fix

* common tests for jump models

* lint fix

* lint fix:2

* full coverage test

* lint fix

* lint fix and remove redundant tests

* lint fix

* std for test modified for gpu tests jump test

* restore merton tests

* split tests

* lint fix

* remove common test files

* isort fix

* avoid test running when imported

* remove double testing

* alternative import for pytest

* lint fix

* lint fix

* full coverage test
---
 pfhedge/instruments/__init__.py               |   1 +
 pfhedge/instruments/primary/kou_jump.py       | 192 ++++++
 pfhedge/stochastic/__init__.py                |   1 +
 pfhedge/stochastic/kou_jump.py                | 192 ++++++
 tests/instruments/primary/test_kou_jump.py    |  17 +
 tests/instruments/primary/test_merton_jump.py |  16 +-
 tests/stochastic/test_kou_jump.py             |  57 ++
 tests/stochastic/test_merton_jump.py          | 547 +++++++++---------
 8 files changed, 746 insertions(+), 277 deletions(-)
 create mode 100644 pfhedge/instruments/primary/kou_jump.py
 create mode 100644 pfhedge/stochastic/kou_jump.py
 create mode 100644 tests/instruments/primary/test_kou_jump.py
 create mode 100644 tests/stochastic/test_kou_jump.py

diff --git a/pfhedge/instruments/__init__.py b/pfhedge/instruments/__init__.py
index 226c927f..f1d4fc0d 100644
--- a/pfhedge/instruments/__init__.py
+++ b/pfhedge/instruments/__init__.py
@@ -15,6 +15,7 @@
 from .primary.brownian import BrownianStock  # NOQA
 from .primary.cir import CIRRate  # NOQA
 from .primary.heston import HestonStock  # NOQA
+from .primary.kou_jump import KouJumpStock  # noqa: F401
 from .primary.local_volatility import LocalVolatilityStock  # NOQA
 from .primary.merton_jump import MertonJumpStock  # NOQA
 from .primary.rough_bergomi import RoughBergomiStock  # NOQA
diff --git a/pfhedge/instruments/primary/kou_jump.py b/pfhedge/instruments/primary/kou_jump.py
new file mode 100644
index 00000000..786bf605
--- /dev/null
+++ b/pfhedge/instruments/primary/kou_jump.py
@@ -0,0 +1,192 @@
+from math import ceil
+from typing import Callable
+from typing import Optional
+from typing import Tuple
+from typing import cast
+
+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_kou_jump
+
+from .base import BasePrimary
+
+
+class KouJumpStock(BasePrimary):
+    r"""A stock of which spot prices follow the Kou's jump diffusion.
+
+    .. seealso::
+        - :func:`pfhedge.stochastic.generate_kou_jump`:
+          The stochastic process.
+
+    Args:
+        sigma (float, default=0.2): The parameter :math:`\sigma`,
+            which stands for the volatility of the spot price.
+        mu (float, default=0.0): The parameter :math:`\mu`,
+            which stands for the drift of the spot price.
+        jump_per_year (float, optional): Jump poisson process annual
+            lambda: Average number of annual jumps. Defaults to 1.0.
+        jump_mean_up (float, optional): Mu for the up jumps:
+            Instaneous value. Defaults to 0.02.
+            This has to be positive and smaller than 1.
+        jump_mean_down (float, optional): Mu for the down jumps:
+            Instaneous value. Defaults to 0.05.
+            This has to be larger than 0.
+        jump_up_prob (float, optional): Given a jump occurs,
+            this is conditional prob for up jump.
+            Down jump occurs with prob 1-jump_up_prob.
+            Has to be in [0,1].
+        cost (float, default=0.0): The transaction cost rate.
+        dt (float, default=1/250): The intervals of the time steps.
+        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.
+            Only to be used for the normal component,
+            jupms uses poisson distribution.
+
+    Buffers:
+        - spot (:class:`torch.Tensor`): The spot prices of the instrument.
+          This attribute is set by a 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.
+
+    Examples:
+        >>> from pfhedge.instruments import KouJumpStock
+        >>>
+        >>> _ = torch.manual_seed(42)
+        >>> stock = KouJumpStock(jump_per_year = 10.0)
+        >>> stock.simulate(n_paths=2, time_horizon=5 / 250)
+        >>> stock.spot
+        tensor([[1.0000, 1.0021, 1.0055, 1.0089, 0.9952, 0.9933],
+                [1.0000, 0.9924, 0.9987, 1.0025, 1.0098, 1.0207]])
+        >>> 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]])
+    """
+
+    def __init__(
+        self,
+        sigma: float = 0.2,
+        mu: float = 0.0,
+        jump_per_year: float = 68.0,
+        jump_mean_up: float = 0.02,
+        jump_mean_down: float = 0.05,
+        jump_up_prob: float = 0.5,
+        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.sigma = sigma
+        self.mu = mu
+        self.jump_per_year = jump_per_year
+        self.jump_mean_up = jump_mean_up
+        self.jump_mean_down = jump_mean_down
+        self.jump_up_prob = jump_up_prob
+        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),)` where
+                :math:`S(0)` is the initial value of the spot price.
+                If ``None`` (default), it uses the default value
+                (See :attr:`default_init_state`).
+                It also accepts a :class:`float` or a :class:`torch.Tensor`.
+        """
+        if init_state is None:
+            init_state = cast(Tuple[float], self.default_init_state)
+
+        spot = generate_kou_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_up=self.jump_mean_up,
+            jump_mean_down=self.jump_mean_down,
+            jump_up_prob=self.jump_up_prob,
+            dt=self.dt,
+            dtype=self.dtype,
+            device=self.device,
+            engine=self.engine,
+        )
+
+        self.register_buffer("spot", spot)
+
+    def extra_repr(self) -> str:
+        params = ["sigma=" + _format_float(self.sigma)]
+        params.append("mu=" + _format_float(self.mu))
+        params.append("cost=" + _format_float(self.cost))
+        params.append("dt=" + _format_float(self.dt))
+        params.append("jump_per_year=" + _format_float(self.jump_per_year))
+        params.append("jump_mean_up=" + _format_float(self.jump_mean_up))
+        params.append("jump_mean_down=" + _format_float(self.jump_mean_down))
+        params.append("jump_up_prob=" + _format_float(self.jump_up_prob))
+        return ", ".join(params)
+
+
+# Assign docstrings so they appear in Sphinx documentation
+_set_docstring(KouJumpStock, "default_init_state", BasePrimary.default_init_state)
+_set_attr_and_docstring(KouJumpStock, "to", BasePrimary.to)
diff --git a/pfhedge/stochastic/__init__.py b/pfhedge/stochastic/__init__.py
index 3861aae0..b6480b6c 100644
--- a/pfhedge/stochastic/__init__.py
+++ b/pfhedge/stochastic/__init__.py
@@ -2,6 +2,7 @@
 from .brownian import generate_geometric_brownian  # NOQA
 from .cir import generate_cir  # NOQA
 from .heston import generate_heston  # NOQA
+from .kou_jump import generate_kou_jump  # noqa: F401
 from .local_volatility import generate_local_volatility_process  # NOQA
 from .merton_jump import generate_merton_jump  # NOQA
 from .random import randn_antithetic  # NOQA
diff --git a/pfhedge/stochastic/kou_jump.py b/pfhedge/stochastic/kou_jump.py
new file mode 100644
index 00000000..ae6bbf3d
--- /dev/null
+++ b/pfhedge/stochastic/kou_jump.py
@@ -0,0 +1,192 @@
+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 ._utils import cast_state
+
+
+def generate_kou_jump(
+    n_paths: int,
+    n_steps: int,
+    init_state: Union[Tuple[TensorOrScalar, ...], TensorOrScalar] = (1.0,),
+    sigma: float = 0.2,
+    mu: float = 0.0,
+    jump_per_year: float = 68.0,
+    jump_mean_up: float = 0.02,
+    jump_mean_down: float = 0.05,
+    jump_up_prob: float = 0.5,
+    dt: float = 1 / 250,
+    dtype: Optional[torch.dtype] = None,
+    device: Optional[torch.device] = None,
+    engine: Callable[..., Tensor] = torch.randn,
+) -> Tensor:
+    r"""Kou's Jump Diffusion Model for stock prices.
+       Assumes number of jumps to be poisson distribution
+       with ASYMMETRIC jump for up and down movement; the
+       lof of these jump follows exponential distribution
+       with mean jump_mean_up and jump_mean_down resp.
+
+       See Glasserman, Paul. Monte Carlo Methods in Financial
+       Engineering. New York: Springer-Verlag, 2004.for details.
+       Copy available at "https://www.bauer.uh.edu/spirrong/
+       Monte_Carlo_Methods_In_Financial_Enginee.pdf"
+
+       Combined with the original paper by Kou:
+       A Jump-Diffusion Model for Option Pricing
+       https://www.columbia.edu/~sk75/MagSci02.pdf
+
+    Args:
+        n_paths (int): The number of simulated paths.
+        n_steps (int): The number of time steps.
+        init_state (tuple[torch.Tensor | float], default=(0.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`.
+            The shape of torch.Tensor must be (1,) or (n_paths,).
+        sigma (float, default=0.2): The parameter :math:`\sigma`,
+            which stands for the volatility of the time series.
+        mu (float, default=0.0): The parameter :math:`\mu`,
+            which stands for the drift of the time series.
+        jump_per_year (float, optional): Jump poisson process annual
+            lambda: Average number of annual jumps. Defaults to 1.0.
+        jump_mean_up (float, optional): Mu for the up jumps:
+            Instaneous value. Defaults to 0.02.
+            This has to be postive and smaller than 1.
+        jump_mean_down (float, optional): Mu for the down jumps:
+            Instaneous value. Defaults to 0.05.
+            This has to be larger than 0.
+        jump_up_prob (float, optional): Given a jump occurs,
+            this is conditional prob for up jump.
+            Down jump occurs with prob 1-jump_up_prob.
+            Has to be in [0,1].
+        dt (float, default=1/250): The intervals of the time steps.
+        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.
+            Only to be used for the normal component,
+            jupms uses poisson 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_kou_jump
+        >>>
+        >>> _ = torch.manual_seed(42)
+        >>> generate_kou_jump(2, 5, jump_per_year = 10.0)
+        tensor([[1.0000, 1.0021, 1.0055, 1.0089, 0.9952],
+                [1.0000, 1.0288, 1.0210, 1.0275, 1.0314]])
+    """
+    if not (0 < jump_mean_up < 1.0):
+        raise ValueError("jump_mean_up must be postive and smaller than 1")
+
+    if not jump_mean_down > 0:
+        raise ValueError("jump_mean_down must be postive")
+
+    if not (0 <= jump_up_prob <= 1.0):
+        raise ValueError("jump prob must be in 0 and 1 incl")
+
+    # change means to rate of exponential distributions
+    jump_eta_up = 1 / jump_mean_up
+    jump_eta_down = 1 / jump_mean_down
+
+    init_state = cast_state(init_state, dtype=dtype, device=device)
+
+    init_value = init_state[0]
+
+    t = dt * torch.arange(n_steps, device=device, dtype=dtype)[None, :]
+    returns = (
+        engine(*(n_paths, n_steps), dtype=dtype, device=device) * math.sqrt(dt) * sigma
+    )
+
+    returns[:, 0] = 0.0
+
+    # Generate jump components
+    poisson = torch.distributions.poisson.Poisson(rate=jump_per_year * dt)
+    n_jumps = poisson.sample((n_paths, n_steps - 1)).to(dtype=dtype, device=device)
+
+    # if n_steps is greater than 1
+    if (n_steps - 1) > 0:
+        # max jumps used to aggregte jump in between dt time
+        max_jumps = int(n_jumps.max())
+        size_paths = torch.Size([n_paths, n_steps - 1, max_jumps])
+
+        # up exp generator
+        up_exp_dist = torch.distributions.exponential.Exponential(rate=jump_eta_up)
+
+        # down exp generator
+        down_exp_dist = torch.distributions.exponential.Exponential(rate=jump_eta_down)
+
+        # up or down generator
+        direction_uni_dist = torch.distributions.uniform.Uniform(0.0, 1.0)
+
+        log_jump = torch.where(
+            direction_uni_dist.sample(size_paths) < jump_up_prob,
+            up_exp_dist.sample(size_paths),
+            -down_exp_dist.sample(size_paths),
+        ).to(returns)
+
+        # for no jump condition
+        log_jump = torch.cat(
+            (torch.zeros(n_paths, n_steps - 1, 1).to(log_jump), log_jump), dim=-1
+        )
+
+        exp_jump_ind = torch.exp(log_jump)
+
+        # filter out jump movements that did not occur in dt time
+        indices_expanded = n_jumps[..., None]
+        k_range = torch.arange(max_jumps + 1).to(returns)
+        mask = k_range > indices_expanded
+        # exp(0) as to no jump after n_jump
+        exp_jump_ind[mask] = 1.0
+
+        # aggregate jumps in time dt--> multiplication of exponent
+        exp_jump = torch.prod(exp_jump_ind, dim=-1)
+
+        # no jump at time 0--> exp(0.0)=1.0
+        exp_jump = torch.cat((torch.ones(n_paths, 1).to(exp_jump), exp_jump), dim=1)
+
+    else:
+        exp_jump = torch.ones(n_paths, 1).to(returns)
+
+    # aggregate jumps upto time t
+    exp_jump_agg = torch.cumprod(exp_jump, dim=-1)
+
+    # jump correction for drift: see the paper
+    m = (
+        (1 - jump_up_prob) * (jump_eta_down / (jump_eta_down + 1))
+        + (jump_up_prob) * (jump_eta_up / (jump_eta_up - 1))
+        - 1
+    )
+
+    prices = (
+        torch.exp(
+            (mu - jump_per_year * m) * t + returns.cumsum(1) - (sigma ** 2) * t / 2
+        )
+        * init_value.view(-1, 1)
+        * exp_jump_agg
+    )
+
+    return prices
diff --git a/tests/instruments/primary/test_kou_jump.py b/tests/instruments/primary/test_kou_jump.py
new file mode 100644
index 00000000..f28b29b8
--- /dev/null
+++ b/tests/instruments/primary/test_kou_jump.py
@@ -0,0 +1,17 @@
+import pytest
+import torch
+
+from pfhedge.instruments import KouJumpStock
+from tests.instruments.primary.test_merton_jump import (
+    TestMertonJumpStock as BaseJumpStockTest,
+)
+
+
+class TestKouJumpStock(BaseJumpStockTest):
+    cls = KouJumpStock
+
+    def test_repr(self):
+        s = KouJumpStock(cost=1e-4)
+        expect = "KouJumpStock(\
+sigma=0.2000, mu=0., cost=1.0000e-04, dt=0.0040, jump_per_year=68., jump_mean_up=0.0200, jump_mean_down=0.0500, jump_up_prob=0.5000)"
+        assert repr(s) == expect
diff --git a/tests/instruments/primary/test_merton_jump.py b/tests/instruments/primary/test_merton_jump.py
index d205c2d7..1bd98810 100644
--- a/tests/instruments/primary/test_merton_jump.py
+++ b/tests/instruments/primary/test_merton_jump.py
@@ -5,14 +5,20 @@
 
 
 class TestMertonJumpStock:
+    cls = MertonJumpStock
+
+    def setup_method(self):
+        self.jump_test_class = self.cls
+
     @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 = self.jump_test_class().to(device)
         s.simulate(n_paths=1000)
 
         assert not s.variance.isnan().any()
+        assert not s.volatility.isnan().any()
 
     @pytest.mark.gpu
     @pytest.mark.parametrize("seed", range(1))
@@ -20,21 +26,23 @@ 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)
+        s = self.jump_test_class(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 = self.jump_test_class(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))
+        assert s.volatility.size() == torch.Size((10, 3))
 
-        s = MertonJumpStock(dt=0.1).to(device)
+        s = self.jump_test_class(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))
+        assert s.volatility.size() == torch.Size((10, 4))
 
     @pytest.mark.gpu
     def test_simulate_shape_gpu(self):
diff --git a/tests/stochastic/test_kou_jump.py b/tests/stochastic/test_kou_jump.py
new file mode 100644
index 00000000..7ddacf4d
--- /dev/null
+++ b/tests/stochastic/test_kou_jump.py
@@ -0,0 +1,57 @@
+from math import sqrt
+
+import pytest
+import torch
+from torch.testing import assert_close
+
+from pfhedge.stochastic import generate_kou_jump
+from pfhedge.stochastic.engine import RandnSobolBoxMuller
+from tests.stochastic.test_merton_jump import (
+    TestGenerateMertonJumpStock as BaseGenerateJumpStockTest,
+)
+
+
+class TestGenerateKouJumpStock(BaseGenerateJumpStockTest):
+    func = staticmethod(generate_kou_jump)
+
+    def test_generate_brownian_mean_no_jump(self, device: str = "cpu"):
+        # kou jump has no std
+        return
+
+    def test_generate_brownian_mean_no_jump_std(self, device: str = "cpu"):
+        # kou jump has no std
+        return
+
+    def test_generate_jump_nosigma2(self, device: str = "cpu"):
+        # kou jump has no std
+        return
+
+    def test_generate_jump_std2(self, device: str = "cpu"):
+        # kou jump has no std
+        return
+
+    # addtional tests for Kou jumo model params
+    def test_kou_jump_mean_up(self):
+        n_paths = 10000
+        n_steps = 250
+        jump_mean_up = 1.1
+        with (
+            pytest.raises(
+                ValueError, match="jump_mean_up must be postive and smaller than 1"
+            )
+        ):
+            generate_kou_jump(n_paths, n_steps, jump_mean_up=jump_mean_up)
+
+    def test_kou_jump_mean_down(self):
+        n_paths = 10000
+        n_steps = 250
+        jump_mean_down = 0.0
+        with pytest.raises(ValueError, match="jump_mean_down must be postive"):
+            generate_kou_jump(n_paths, n_steps, jump_mean_down=jump_mean_down)
+
+    def test_kou_up_jump_prob(self):
+        n_paths = 10000
+        n_steps = 250
+        jump_up_prob = 1.1
+        with pytest.raises(ValueError, match="jump prob must be in 0 and 1 incl"):
+            generate_kou_jump(n_paths, n_steps, jump_up_prob=jump_up_prob)
diff --git a/tests/stochastic/test_merton_jump.py b/tests/stochastic/test_merton_jump.py
index f3c00c3d..1e7c8adc 100644
--- a/tests/stochastic/test_merton_jump.py
+++ b/tests/stochastic/test_merton_jump.py
@@ -8,276 +8,277 @@
 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=68.2,  # default value
-        jump_std=0.0,
-        jump_mean=0.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.4 * 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].mean()
-    expect = torch.ones_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")
+class TestGenerateMertonJumpStock:
+    func = staticmethod(generate_merton_jump)
+
+    def setup_method(self):
+        self.jump_test_func = self.func
+
+    def test_generate_brownian_mean_no_jump(self, device: str = "cpu"):
+        torch.manual_seed(42)
+        n_paths = 10000
+        n_steps = 250
+
+        output = self.jump_test_func(
+            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(self):
+        self.test_generate_brownian_mean_no_jump(device="cuda")
+
+    def test_generate_brownian_mean_no_jump1(self, device: str = "cpu"):
+        torch.manual_seed(42)
+        n_paths = 10000
+        n_steps = 250
+
+        output = self.jump_test_func(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(self):
+        self.test_generate_brownian_mean_no_jump1(device="cuda")
+
+    def test_generate_brownian_mean_no_jump_std(self, device: str = "cpu"):
+        torch.manual_seed(42)
+        n_paths = 10000
+        n_steps = 250
+
+        output = self.jump_test_func(
+            n_paths,
+            n_steps,
+            jump_per_year=68.2,  # default value
+            jump_std=0.0,
+            jump_mean=0.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.5 * 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(self):
+        self.test_generate_brownian_mean_no_jump_std(device="cuda")
+
+    def test_generate_brownian_mean(self, device: str = "cpu"):
+        torch.manual_seed(42)
+        n_paths = 10000
+        n_steps = 250
+
+        output = self.jump_test_func(
+            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(self):
+        self.test_generate_brownian_mean(device="cuda")
+
+    def test_generate_jump_nosigma(self, device: str = "cpu"):
+        torch.manual_seed(42)
+        n_steps = 250
+
+        result = self.jump_test_func(
+            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 = self.jump_test_func(
+            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_jump_nosigma_gpu(self):
+        self.test_generate_jump_nosigma(device="cpu")
+
+    def test_generate_jump_nosigma2(self, device: str = "cpu"):
+        torch.manual_seed(42)
+        n_steps = 250
+
+        result = self.jump_test_func(
+            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 = self.jump_test_func(
+            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_jump_nosigma2_gpu(self):
+        self.test_generate_jump_nosigma2(device="cuda")
+
+    def test_generate_jump_std(self, device: str = "cpu"):
+        torch.manual_seed(42)
+        n_paths = 10000
+        n_steps = 250
+
+        output = self.jump_test_func(
+            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_jump_std_gpu(self):
+        self.test_generate_jump_std(device="cuda")
+
+    def test_generate_jump_std2(self, device: str = "cpu"):
+        torch.manual_seed(42)
+        n_paths = 10000
+        n_steps = 250
+
+        output = self.jump_test_func(
+            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_jump_std2_gpu(self):
+        self.test_generate_jump_std2(device="cuda")
+
+    def test_generate_jump_mean_init_state(self, device: str = "cpu"):
+        torch.manual_seed(42)
+        n_paths = 10000
+        n_steps = 250
+
+        output = self.jump_test_func(
+            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 = self.jump_test_func(
+            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 = self.jump_test_func(
+            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_jump_mean_init_state_gpu(self):
+        self.test_generate_jump_mean_init_state(device="cuda")
+
+    def test_generate_jump_mean_mu(self, device: str = "cpu"):
+        torch.manual_seed(42)
+        n_paths = 10000
+        n_steps = 250
+        dt = 1 / 250
+        mu = 0.1
+
+        output = self.jump_test_func(
+            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_jump_mean_mu_gpu(self):
+        self.test_generate_jump_mean_mu(device="cuda")
+
+    def test_generate_jump_dtype(self, device: str = "cpu"):
+        torch.manual_seed(42)
+
+        output = self.jump_test_func(
+            1, 1, dtype=torch.float32, device=torch.device(device)
+        )
+        assert output.dtype == torch.float32
+
+        output = self.jump_test_func(
+            1, 1, dtype=torch.float64, device=torch.device(device)
+        )
+        assert output.dtype == torch.float64
+
+    @pytest.mark.gpu
+    def test_generate_jump_dtype_gpu(self):
+        self.test_generate_jump_dtype(device="cuda")
+
+    def test_generate_jump_sobol_mean(self, device: str = "cpu"):
+        n_paths = 10000
+        n_steps = 250
+
+        engine = RandnSobolBoxMuller(seed=42, scramble=True)
+        output = self.jump_test_func(
+            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].mean()
+        expect = torch.ones_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_jump_sobol_mean_gpu(self):
+        self.test_generate_jump_sobol_mean(device="cuda")