test and lit fix

pfhedge/instruments/primary/kou_jump.py

@@ -61,8 +61,8 @@ class KouJumpStock(BasePrimary):
         >>> stock = KouJumpStock()
         >>> stock.simulate(n_paths=2, time_horizon=5 / 250)
         >>> stock.spot
-        tensor([[1.0000, 1.0101, 1.0137, 1.0144, 1.0211, 1.0180],
-                [1.0000, 1.0067, 1.0065, 1.0158, 1.0175, 1.0287]])
+        tensor([[1.0000, 1.0018, 1.0084, 1.0150, 1.0044, 1.0056],
+                [1.0000, 0.9956, 1.0050, 1.0121, 1.0227, 1.0369]])
 
         Using custom ``dtype`` and ``device``.
 
@@ -149,8 +149,8 @@ def simulate(
             >>> stock = KouJumpStock()
             >>> stock.simulate(n_paths=2, time_horizon=5 / 250, init_state=(2.0,))
             >>> stock.spot
-            tensor([[2.0000, 2.0032, 2.0091, 2.0149, 1.9865, 1.9817],
-                    [2.0000, 1.9839, 1.9954, 2.0022, 2.0157, 2.0364]])
+            tensor([[2.0000, 2.0036, 2.0169, 2.0301, 2.0087, 2.0113],
+                    [2.0000, 1.9911, 2.0100, 2.0242, 2.0453, 2.0738]])
         """
         if init_state is None:
             init_state = cast(Tuple[float], self.default_init_state)

pfhedge/stochastic/kou_jump.py

@@ -93,8 +93,8 @@ def generate_kou_jump(
         >>>
         >>> _ = torch.manual_seed(42)
         >>> generate_kou_jump(2, 5)
-        tensor([[1.0000, 1.0053, 1.0119, 0.9272, 0.9174],
-                [1.0000, 1.0321, 1.0275, 1.0373, 1.0446]])
+        tensor([[1.0000, 1.0053, 1.0119, 0.9271, 0.9174],
+                [1.0000, 1.0321, 1.0275, 1.0372, 1.0445]])
     """
     assert jump_eta_up > 1.0, "jump_eta_up must be larger than 1.0"
     assert jump_eta_down > 0.0, "jump_eta_down must be larger than 0.0"
@@ -116,7 +116,7 @@ def generate_kou_jump(
         returns
     )
     # if n_steps is greater than 1
-    if (n_steps-1 ) > 0:
+    if (n_steps - 1) > 0:
         # max jumps used to aggregte jump in between dt time
         max_jumps = int(n_jumps.max())
         size_paths = (n_paths, n_steps - 1, max_jumps)
@@ -137,7 +137,7 @@ def generate_kou_jump(
         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
@@ -167,7 +167,9 @@ def generate_kou_jump(
     )
 
     prices = (
-        torch.exp((mu - jump_per_year * m) * t + returns.cumsum(1) - (sigma ** 2) * t / 2)
+        torch.exp(
+            (mu - jump_per_year * m) * t + returns.cumsum(1) - (sigma ** 2) * t / 2
+        )
         * init_value.view(-1, 1)
         * exp_jump_agg
     )

tests/stochastic/test_kou_jump.py

@@ -154,14 +154,10 @@ def test_generate_kou_jump_mean_mu_gpu():
 def test_generate_kou_jump_dtype(device: str = "cpu"):
     torch.manual_seed(42)
 
-    output = generate_kou_jump(
-        1, 1, dtype=torch.float32, device=torch.device(device)
-    )
+    output = generate_kou_jump(1, 1, dtype=torch.float32, device=torch.device(device))
     assert output.dtype == torch.float32
 
-    output = generate_kou_jump(
-        1, 1, dtype=torch.float64, device=torch.device(device)
-    )
+    output = generate_kou_jump(1, 1, dtype=torch.float64, device=torch.device(device))
     assert output.dtype == torch.float64