From c69524bb36af28b0776dd214b4466031394d9ae3 Mon Sep 17 00:00:00 2001
From: samueledelia <delia.samuele00@gmail.com>
Date: Mon, 19 Feb 2024 17:54:22 +0100
Subject: [PATCH 1/6] Update MultiGBM, TestMultiGBM

---
 .idea/misc.xml                |  2 +-
 .idea/py-exopricer.iml        |  2 +-
 jaxfin/models/gbm/__init__.py |  4 +-
 jaxfin/models/gbm/gbm.py      | 80 +++++++++++++++++++++++++++++++++++
 tests/models/test_gbm.py      | 17 +++++++-
 5 files changed, 100 insertions(+), 5 deletions(-)

diff --git a/.idea/misc.xml b/.idea/misc.xml
index 870738f..034c334 100644
--- a/.idea/misc.xml
+++ b/.idea/misc.xml
@@ -4,5 +4,5 @@
     <option name="pathToExecutable" value="C:\Users\paolo\AppData\Local\Programs\Python\Python310\Scripts\black.exe" />
     <option name="sdkName" value="Python 3.10 (py-exopricer)" />
   </component>
-  <component name="ProjectRootManager" version="2" project-jdk-name="Python 3.10 (py-exopricer)" project-jdk-type="Python SDK" />
+  <component name="ProjectRootManager" version="2" project-jdk-name="Python 3.9 (py-exopricer)" project-jdk-type="Python SDK" />
 </project>
\ No newline at end of file
diff --git a/.idea/py-exopricer.iml b/.idea/py-exopricer.iml
index 74d515a..fab6b99 100644
--- a/.idea/py-exopricer.iml
+++ b/.idea/py-exopricer.iml
@@ -4,7 +4,7 @@
     <content url="file://$MODULE_DIR$">
       <excludeFolder url="file://$MODULE_DIR$/venv" />
     </content>
-    <orderEntry type="inheritedJdk" />
+    <orderEntry type="jdk" jdkName="Python 3.9 (py-exopricer)" jdkType="Python SDK" />
     <orderEntry type="sourceFolder" forTests="false" />
   </component>
 </module>
\ No newline at end of file
diff --git a/jaxfin/models/gbm/__init__.py b/jaxfin/models/gbm/__init__.py
index f50bea8..c0d9614 100644
--- a/jaxfin/models/gbm/__init__.py
+++ b/jaxfin/models/gbm/__init__.py
@@ -1,6 +1,6 @@
 """
 Geometric Brownian motion module
 """
-from jaxfin.models.gbm.gbm import UnivGeometricBrownianMotion
+from jaxfin.models.gbm.gbm import UnivGeometricBrownianMotion, MultiGeometricBrownianMotion
 
-__all__ = ["UnivGeometricBrownianMotion"]
+__all__ = ["UnivGeometricBrownianMotion", "MultiGeometricBrownianMotion"]
diff --git a/jaxfin/models/gbm/gbm.py b/jaxfin/models/gbm/gbm.py
index d5d8960..4b4eda4 100644
--- a/jaxfin/models/gbm/gbm.py
+++ b/jaxfin/models/gbm/gbm.py
@@ -82,3 +82,83 @@ def simulate_paths(
         Xt = jnp.vstack([jnp.ones(n_sim), Xt])
 
         return self._s0 * Xt.cumprod(axis=0)
+
+class MultiGeometricBrownianMotion:
+        """
+        Geometric Brownian Motion
+
+        Represent a d-dimensional GBM
+
+        # Example usage:
+        params = {
+            's0' : [10, 12],
+            'dtype' : jnp.float32,
+            'mean' : 0.1,
+            'sigma': 0.3
+        }
+
+        gmb_process = GeometricBrownianMotion(**params)
+        paths = gmb_process.simulate_paths(maturity=1.0, n=100, n_sim=100)
+        """
+
+        def __init__(self, s0, mean, cov, dtype):
+            if dtype is None:
+                raise ValueError("dtype must not be None")
+
+            self._dtype = dtype
+            self._s0 = jnp.asarray(s0, dtype=dtype)
+            self._mean = jnp.asarray(mean, dtype=dtype)
+            self._cov = jnp.asarray(cov, dtype=dtype)
+
+        @property
+        def mean(self):
+            """
+            :return: Returns the mean of the GBM
+            """
+            return self._mean
+
+        @property
+        def cov(self):
+            """
+            :return: Returns the standard deviation of the GBM
+            """
+            return self._cov
+
+        @property
+        def s0(self):
+            """
+            :return: Returns the initial value of the GBM
+            """
+            return self._s0
+
+        @property
+        def dtype(self):
+            """
+            :return: Returns the underlying dtype of the GBM
+            """
+            return self._dtype
+
+        def simulate_paths(
+                self, seed: int, maturity, n: int
+        ) -> jax.Array:
+            """
+            Simulate a sample of paths from the GBM
+
+            :param maturity: time in years
+            :param n: (int): number of steps
+            :return: (jax.Array): Array containing the sample paths
+            """
+            key = random.PRNGKey(seed)
+
+            n_assets = self._s0.size[0]
+            choleskyMatrix = jnp.linalg.cholesky(self._cov)
+            e = jnp.random.normal(size=(n_assets, n))
+            correlated_returns = jnp.dot(choleskyMatrix, e)
+
+            # Initialize asset prices
+            s = jnp.ones((n_assets, n)) * self._s0
+            for t in range(1, n):
+                drift = self._mean * s[:, t - 1]
+                diffusion = jnp.diag(self._cov) * s[:, t - 1] * correlated_returns[:, t]
+                s[:, t] = s[:, t - 1] * jnp.exp(drift + diffusion)
+            return s
diff --git a/tests/models/test_gbm.py b/tests/models/test_gbm.py
index e8ae4ba..98c94f1 100644
--- a/tests/models/test_gbm.py
+++ b/tests/models/test_gbm.py
@@ -1,6 +1,6 @@
 import jax.numpy as jnp
 
-from jaxfin.models.gbm import UnivGeometricBrownianMotion
+from jaxfin.models.gbm import UnivGeometricBrownianMotion, MultiGeometricBrownianMotion
 
 
 class TestUnivGBM:
@@ -16,3 +16,18 @@ def test_init(self):
         assert gbm.sigma == sigma
         assert gbm.dtype == dtype
         assert gbm.s0 == s0
+
+
+class TestMultiGBM:
+
+    def test_init(self):
+        s0 = jnp.array([10, 12])
+        mean = jnp.array([0.1, 0.0])
+        cov = jnp.array([[0.3, 0.1], [0.1, 0.5]])
+        dtype = jnp.float32
+        gbm = MultiGeometricBrownianMotion(s0, mean, cov, dtype)
+
+        assert jnp.array_equal(gbm.mean, mean)
+        assert jnp.array_equal(gbm.cov, cov)
+        assert gbm.dtype == dtype
+        assert jnp.array_equal(gbm.s0, s0)
\ No newline at end of file

From 11a2f9d1c173c5a651adfd0366bf6d45dcb96042 Mon Sep 17 00:00:00 2001
From: samueledelia <delia.samuele00@gmail.com>
Date: Mon, 19 Feb 2024 22:42:53 +0100
Subject: [PATCH 2/6] Update MultiGBM, test_sample_path

---
 jaxfin/models/gbm/gbm.py | 58 +++++++++++++++++++++++++++++-----------
 tests/models/test_gbm.py | 22 ++++++++++++---
 2 files changed, 60 insertions(+), 20 deletions(-)

diff --git a/jaxfin/models/gbm/gbm.py b/jaxfin/models/gbm/gbm.py
index 4b4eda4..0f9be45 100644
--- a/jaxfin/models/gbm/gbm.py
+++ b/jaxfin/models/gbm/gbm.py
@@ -94,21 +94,28 @@ class MultiGeometricBrownianMotion:
             's0' : [10, 12],
             'dtype' : jnp.float32,
             'mean' : 0.1,
-            'sigma': 0.3
+            'cov':  [[0.3, 0.1], [0.1, 0.5]]
         }
 
         gmb_process = GeometricBrownianMotion(**params)
         paths = gmb_process.simulate_paths(maturity=1.0, n=100, n_sim=100)
         """
 
-        def __init__(self, s0, mean, cov, dtype):
+        def __init__(self, s0, mean, sigma, corr, dtype):
             if dtype is None:
                 raise ValueError("dtype must not be None")
 
+            if not check_symmetric(corr, 1e-8):
+                raise ValueError("Correlation matrix must be symmetric")
+
+            if not jnp.array_equal(jnp.diag(corr), jnp.ones(corr.shape[0])):
+                raise ValueError("Correlation matrix must have ones as diagonal elements")
+
             self._dtype = dtype
             self._s0 = jnp.asarray(s0, dtype=dtype)
             self._mean = jnp.asarray(mean, dtype=dtype)
-            self._cov = jnp.asarray(cov, dtype=dtype)
+            self._sigma = jnp.asarray(sigma, dtype=dtype)
+            self._corr = jnp.asarray(corr, dtype=dtype)
 
         @property
         def mean(self):
@@ -117,12 +124,18 @@ def mean(self):
             """
             return self._mean
 
-        @property
-        def cov(self):
+
+        def sigma(self):
             """
             :return: Returns the standard deviation of the GBM
             """
-            return self._cov
+            return self._sigma
+        @property
+        def corr(self):
+            """
+            :return: Returns the correlation matrix of the Weiner processes
+            """
+            return self._corr
 
         @property
         def s0(self):
@@ -138,6 +151,7 @@ def dtype(self):
             """
             return self._dtype
 
+
         def simulate_paths(
                 self, seed: int, maturity, n: int
         ) -> jax.Array:
@@ -150,15 +164,27 @@ def simulate_paths(
             """
             key = random.PRNGKey(seed)
 
+            dt = maturity/n
             n_assets = self._s0.size[0]
-            choleskyMatrix = jnp.linalg.cholesky(self._cov)
+
+            choleskyMatrix = jnp.linalg.cholesky(self._corr)
             e = jnp.random.normal(size=(n_assets, n))
-            correlated_returns = jnp.dot(choleskyMatrix, e)
-
-            # Initialize asset prices
-            s = jnp.ones((n_assets, n)) * self._s0
-            for t in range(1, n):
-                drift = self._mean * s[:, t - 1]
-                diffusion = jnp.diag(self._cov) * s[:, t - 1] * correlated_returns[:, t]
-                s[:, t] = s[:, t - 1] * jnp.exp(drift + diffusion)
-            return s
+            stochastic_increment = jnp.dot(choleskyMatrix, e)
+
+            log_increments = ((self._mean - self._sigma**2 / 2) * dt
+                      + jnp.sqrt(dt) * self._sigma
+                      * stochastic_increment)
+
+            once = jnp.ones([n, n], dtype=self._dtype)
+            lower_triangular = jnp.tril(once)
+            cumsum = jnp.dot(lower_triangular,
+                                      jnp.transpose(log_increments))
+            cumsum = jnp.transpose(cumsum, [1, 2, 0])
+            samples = self._s0 * jnp.math.exp(cumsum)
+            return samples
+
+
+
+
+def check_symmetric(a, tol=1e-8):
+    return jnp.all(jnp.abs(a-a.T) < tol)
\ No newline at end of file
diff --git a/tests/models/test_gbm.py b/tests/models/test_gbm.py
index 98c94f1..116e4e8 100644
--- a/tests/models/test_gbm.py
+++ b/tests/models/test_gbm.py
@@ -23,11 +23,25 @@ class TestMultiGBM:
     def test_init(self):
         s0 = jnp.array([10, 12])
         mean = jnp.array([0.1, 0.0])
-        cov = jnp.array([[0.3, 0.1], [0.1, 0.5]])
+        sigma = jnp.array([0.3, 0.5])
+        corr = jnp.array([[1, 0.1], [0.1, 1]])
         dtype = jnp.float32
-        gbm = MultiGeometricBrownianMotion(s0, mean, cov, dtype)
+        gbm = MultiGeometricBrownianMotion(s0, mean, sigma, corr, dtype)
 
         assert jnp.array_equal(gbm.mean, mean)
-        assert jnp.array_equal(gbm.cov, cov)
+        assert jnp.array_equal(gbm.sigma, sigma)
+        assert jnp.array_equal(gbm.corr, corr)
         assert gbm.dtype == dtype
-        assert jnp.array_equal(gbm.s0, s0)
\ No newline at end of file
+        assert jnp.array_equal(gbm.s0, s0)
+
+
+    def test_sample_path(self):
+        s0 = jnp.array([10, 12])
+        mean = jnp.array([0.1, 0.0])
+        sigma = jnp.array([0.3, 0.5])
+        corr = jnp.array([[1, 0.1], [0.1, 1]])
+        dtype = jnp.float32
+        gbm = MultiGeometricBrownianMotion(s0, mean, sigma, corr, dtype)
+        sample_path = gbm.simulate_paths(0, 1, 52)
+
+        assert True

From d4741b0895f7d54d80787841d45c39197b30a628 Mon Sep 17 00:00:00 2001
From: samueledelia <delia.samuele00@gmail.com>
Date: Mon, 19 Feb 2024 22:48:39 +0100
Subject: [PATCH 3/6] Reformatted code

---
 jaxfin/models/gbm/__init__.py |   5 +-
 jaxfin/models/gbm/gbm.py      | 173 +++++++++++++++++-----------------
 jaxfin/price_engine/common.py |   2 +-
 jaxfin/price_engine/math.py   |   2 +-
 jaxfin/price_engine/utils.py  |   2 +-
 5 files changed, 93 insertions(+), 91 deletions(-)

diff --git a/jaxfin/models/gbm/__init__.py b/jaxfin/models/gbm/__init__.py
index c0d9614..18bff58 100644
--- a/jaxfin/models/gbm/__init__.py
+++ b/jaxfin/models/gbm/__init__.py
@@ -1,6 +1,9 @@
 """
 Geometric Brownian motion module
 """
-from jaxfin.models.gbm.gbm import UnivGeometricBrownianMotion, MultiGeometricBrownianMotion
+from jaxfin.models.gbm.gbm import (
+    UnivGeometricBrownianMotion,
+    MultiGeometricBrownianMotion,
+)
 
 __all__ = ["UnivGeometricBrownianMotion", "MultiGeometricBrownianMotion"]
diff --git a/jaxfin/models/gbm/gbm.py b/jaxfin/models/gbm/gbm.py
index 0f9be45..7d79cb9 100644
--- a/jaxfin/models/gbm/gbm.py
+++ b/jaxfin/models/gbm/gbm.py
@@ -59,9 +59,7 @@ def dtype(self):
         """
         return self._dtype
 
-    def simulate_paths(
-        self, seed: int, maturity, n: int, n_sim: int
-    ) -> jax.Array:
+    def simulate_paths(self, seed: int, maturity, n: int, n_sim: int) -> jax.Array:
         """
         Simulate a sample of paths from the GBM
 
@@ -75,7 +73,7 @@ def simulate_paths(
         dt = maturity / n
 
         Xt = jnp.exp(
-            (self._mean - self._sigma ** 2 / 2) * dt
+            (self._mean - self._sigma**2 / 2) * dt
             + self._sigma * jnp.sqrt(dt) * random.normal(key, shape=(n_sim, n - 1)).T
         )
 
@@ -83,108 +81,109 @@ def simulate_paths(
 
         return self._s0 * Xt.cumprod(axis=0)
 
-class MultiGeometricBrownianMotion:
-        """
-        Geometric Brownian Motion
-
-        Represent a d-dimensional GBM
-
-        # Example usage:
-        params = {
-            's0' : [10, 12],
-            'dtype' : jnp.float32,
-            'mean' : 0.1,
-            'cov':  [[0.3, 0.1], [0.1, 0.5]]
-        }
 
-        gmb_process = GeometricBrownianMotion(**params)
-        paths = gmb_process.simulate_paths(maturity=1.0, n=100, n_sim=100)
-        """
+class MultiGeometricBrownianMotion:
+    """
+    Geometric Brownian Motion
 
-        def __init__(self, s0, mean, sigma, corr, dtype):
-            if dtype is None:
-                raise ValueError("dtype must not be None")
+    Represent a d-dimensional GBM
 
-            if not check_symmetric(corr, 1e-8):
-                raise ValueError("Correlation matrix must be symmetric")
+    # Example usage:
+    params = {
+        's0' : [10, 12],
+        'dtype' : jnp.float32,
+        'mean' : 0.1,
+        'cov':  [[0.3, 0.1], [0.1, 0.5]]
+    }
 
-            if not jnp.array_equal(jnp.diag(corr), jnp.ones(corr.shape[0])):
-                raise ValueError("Correlation matrix must have ones as diagonal elements")
+    gmb_process = GeometricBrownianMotion(**params)
+    paths = gmb_process.simulate_paths(maturity=1.0, n=100, n_sim=100)
+    """
 
-            self._dtype = dtype
-            self._s0 = jnp.asarray(s0, dtype=dtype)
-            self._mean = jnp.asarray(mean, dtype=dtype)
-            self._sigma = jnp.asarray(sigma, dtype=dtype)
-            self._corr = jnp.asarray(corr, dtype=dtype)
+    def __init__(self, s0, mean, sigma, corr, dtype):
+        if dtype is None:
+            raise ValueError("dtype must not be None")
 
-        @property
-        def mean(self):
-            """
-            :return: Returns the mean of the GBM
-            """
-            return self._mean
+        if not check_symmetric(corr, 1e-8):
+            raise ValueError("Correlation matrix must be symmetric")
 
+        if not jnp.array_equal(jnp.diag(corr), jnp.ones(corr.shape[0])):
+            raise ValueError("Correlation matrix must have ones as diagonal elements")
 
-        def sigma(self):
-            """
-            :return: Returns the standard deviation of the GBM
-            """
-            return self._sigma
-        @property
-        def corr(self):
-            """
-            :return: Returns the correlation matrix of the Weiner processes
-            """
-            return self._corr
+        self._dtype = dtype
+        self._s0 = jnp.asarray(s0, dtype=dtype)
+        self._mean = jnp.asarray(mean, dtype=dtype)
+        self._sigma = jnp.asarray(sigma, dtype=dtype)
+        self._corr = jnp.asarray(corr, dtype=dtype)
 
-        @property
-        def s0(self):
-            """
-            :return: Returns the initial value of the GBM
-            """
-            return self._s0
+    @property
+    def mean(self):
+        """
+        :return: Returns the mean of the GBM
+        """
+        return self._mean
 
-        @property
-        def dtype(self):
-            """
-            :return: Returns the underlying dtype of the GBM
-            """
-            return self._dtype
+    def sigma(self):
+        """
+        :return: Returns the standard deviation of the GBM
+        """
+        return self._sigma
 
+    @property
+    def corr(self):
+        """
+        :return: Returns the correlation matrix of the Weiner processes
+        """
+        return self._corr
 
-        def simulate_paths(
-                self, seed: int, maturity, n: int
-        ) -> jax.Array:
-            """
-            Simulate a sample of paths from the GBM
+    @property
+    def s0(self):
+        """
+        :return: Returns the initial value of the GBM
+        """
+        return self._s0
 
-            :param maturity: time in years
-            :param n: (int): number of steps
-            :return: (jax.Array): Array containing the sample paths
-            """
-            key = random.PRNGKey(seed)
+    @property
+    def dtype(self):
+        """
+        :return: Returns the underlying dtype of the GBM
+        """
+        return self._dtype
 
-            dt = maturity/n
-            n_assets = self._s0.size[0]
+    def simulate_paths(self, seed: int, maturity, n: int) -> jax.Array:
+        """
+        Simulate a sample of paths from the GBM
 
-            choleskyMatrix = jnp.linalg.cholesky(self._corr)
-            e = jnp.random.normal(size=(n_assets, n))
-            stochastic_increment = jnp.dot(choleskyMatrix, e)
+        :param maturity: time in years
+        :param n: (int): number of steps
+        :return: (jax.Array): Array containing the sample paths
+        """
+        key = random.PRNGKey(seed)
 
-            log_increments = ((self._mean - self._sigma**2 / 2) * dt
-                      + jnp.sqrt(dt) * self._sigma
-                      * stochastic_increment)
+        dt = maturity / n
+        n_assets = self._s0.size[0]
 
-            once = jnp.ones([n, n], dtype=self._dtype)
-            lower_triangular = jnp.tril(once)
-            cumsum = jnp.dot(lower_triangular,
-                                      jnp.transpose(log_increments))
-            cumsum = jnp.transpose(cumsum, [1, 2, 0])
-            samples = self._s0 * jnp.math.exp(cumsum)
-            return samples
+        choleskyMatrix = jnp.linalg.cholesky(self._corr)
+        e = random.normal(key, shape=(n_assets, n))
+        stochastic_increment = jnp.dot(choleskyMatrix, e)
 
+        log_increments = (self._mean - self._sigma**2 / 2) * dt + jnp.sqrt(
+            dt
+        ) * self._sigma * stochastic_increment
 
+        once = jnp.ones([n, n], dtype=self._dtype)
+        lower_triangular = jnp.tril(once)
+        cumsum = jnp.dot(lower_triangular, jnp.transpose(log_increments))
+        cumsum = jnp.transpose(cumsum, [1, 2, 0])
+        samples = self._s0 * jnp.math.exp(cumsum)
+        return samples
 
 
 def check_symmetric(a, tol=1e-8):
-    return jnp.all(jnp.abs(a-a.T) < tol)
\ No newline at end of file
+    """
+
+    :param a:
+    :param tol:
+    :return:
+    """
+    return jnp.all(jnp.abs(a - a.T) < tol)
diff --git a/jaxfin/price_engine/common.py b/jaxfin/price_engine/common.py
index 4029b90..13477f8 100644
--- a/jaxfin/price_engine/common.py
+++ b/jaxfin/price_engine/common.py
@@ -46,7 +46,7 @@ def compute_undiscounted_call_prices(spots, strikes, expires, vols, discount_rat
 def compute_discounted_call_prices(spots, strikes, expires, vols, discount_rates):
     """
     Compute the discounted call option prices
-    
+
     :param spots: Current spot price of the underlying
     :param strikes: Strike price of the option
     :param expires: Time to expiration of the option
diff --git a/jaxfin/price_engine/math.py b/jaxfin/price_engine/math.py
index 3ba1faa..958824d 100644
--- a/jaxfin/price_engine/math.py
+++ b/jaxfin/price_engine/math.py
@@ -45,7 +45,7 @@ def cum_normal(x):
 def density_normal(x):
     """
     Normal distribution function
-    
+
     :param x: Input value
     :return: Normal distribution value
     """
diff --git a/jaxfin/price_engine/utils.py b/jaxfin/price_engine/utils.py
index e36370b..50e93ed 100644
--- a/jaxfin/price_engine/utils.py
+++ b/jaxfin/price_engine/utils.py
@@ -25,7 +25,7 @@ def cast_arrays(array: List[jax.Array], dtype):
 def check_shape(*args):
     """
     Checks if the shapes of the input arrays are the same
-    
+
     :param args: List of arrays
     :return: True if the shapes are the same, False otherwise
     """

From efe1606fecae613ca8bc1327c56c36ecc728b0ef Mon Sep 17 00:00:00 2001
From: Paolo D'Elia <paolo.delia99@gmail.com>
Date: Tue, 20 Feb 2024 15:06:30 +0100
Subject: [PATCH 4/6] Implemented multivariate gbm

---
 jaxfin/models/gbm/gbm.py | 46 ++++++++++++++++++++++++++--------------
 tests/models/test_gbm.py | 17 +++++++++++++--
 2 files changed, 45 insertions(+), 18 deletions(-)

diff --git a/jaxfin/models/gbm/gbm.py b/jaxfin/models/gbm/gbm.py
index 7d79cb9..43b3770 100644
--- a/jaxfin/models/gbm/gbm.py
+++ b/jaxfin/models/gbm/gbm.py
@@ -104,7 +104,7 @@ def __init__(self, s0, mean, sigma, corr, dtype):
         if dtype is None:
             raise ValueError("dtype must not be None")
 
-        if not check_symmetric(corr, 1e-8):
+        if not _check_symmetric(corr, 1e-8):
             raise ValueError("Correlation matrix must be symmetric")
 
         if not jnp.array_equal(jnp.diag(corr), jnp.ones(corr.shape[0])):
@@ -115,6 +115,7 @@ def __init__(self, s0, mean, sigma, corr, dtype):
         self._mean = jnp.asarray(mean, dtype=dtype)
         self._sigma = jnp.asarray(sigma, dtype=dtype)
         self._corr = jnp.asarray(corr, dtype=dtype)
+        self._dim = self._s0.shape[0]
 
     @property
     def mean(self):
@@ -123,6 +124,7 @@ def mean(self):
         """
         return self._mean
 
+    @property
     def sigma(self):
         """
         :return: Returns the standard deviation of the GBM
@@ -150,7 +152,14 @@ def dtype(self):
         """
         return self._dtype
 
-    def simulate_paths(self, seed: int, maturity, n: int) -> jax.Array:
+    @property
+    def dimension(self):
+        """
+        :return: Returns the dimension of the GBM
+        """
+        return self._dim
+
+    def simulate_paths(self, seed: int, maturity, n: int, n_sim: int) -> jax.Array:
         """
         Simulate a sample of paths from the GBM
 
@@ -161,29 +170,34 @@ def simulate_paths(self, seed: int, maturity, n: int) -> jax.Array:
         key = random.PRNGKey(seed)
 
         dt = maturity / n
-        n_assets = self._s0.size[0]
 
-        choleskyMatrix = jnp.linalg.cholesky(self._corr)
-        e = random.normal(key, shape=(n_assets, n))
-        stochastic_increment = jnp.dot(choleskyMatrix, e)
+        normal_draw = random.normal(key, shape=(n_sim, n * self._dim))
+        normal_draw = jnp.reshape(normal_draw, (n_sim, n, self._dim)).transpose(
+            (1, 0, 2)
+        )
+
+        cholesky_matrix = jnp.linalg.cholesky(self._corr)
+
+        stochastic_increments = normal_draw @ cholesky_matrix
 
         log_increments = (self._mean - self._sigma**2 / 2) * dt + jnp.sqrt(
             dt
-        ) * self._sigma * stochastic_increment
+        ) * self._sigma * stochastic_increments
 
         once = jnp.ones([n, n], dtype=self._dtype)
-        lower_triangular = jnp.tril(once)
-        cumsum = jnp.dot(lower_triangular, jnp.transpose(log_increments))
-        cumsum = jnp.transpose(cumsum, [1, 2, 0])
-        samples = self._s0 * jnp.math.exp(cumsum)
-        return samples
+        lower_triangular = jnp.tril(once, k=-1)
+        cumsum = log_increments.transpose() @ lower_triangular
+        cumsum = cumsum.transpose((1, 2, 0))
+        samples = self._s0 * jnp.exp(cumsum)
+        return samples.transpose(1, 0, 2)[::-1, :, :]
 
 
-def check_symmetric(a, tol=1e-8):
+def _check_symmetric(a, tol=1e-8):
     """
+    Check if a matrix is symmetric
 
-    :param a:
-    :param tol:
-    :return:
+    :param a: (jax.Array): Matrix to check
+    :param tol: (float): Tolerance for the check
+    :return: (bool): True if the matrix is symmetric
     """
     return jnp.all(jnp.abs(a - a.T) < tol)
diff --git a/tests/models/test_gbm.py b/tests/models/test_gbm.py
index 116e4e8..fc325d9 100644
--- a/tests/models/test_gbm.py
+++ b/tests/models/test_gbm.py
@@ -2,6 +2,7 @@
 
 from jaxfin.models.gbm import UnivGeometricBrownianMotion, MultiGeometricBrownianMotion
 
+SEED: int = 42
 
 class TestUnivGBM:
 
@@ -17,6 +18,17 @@ def test_init(self):
         assert gbm.dtype == dtype
         assert gbm.s0 == s0
 
+    def test_sim_paths_shape(self):
+        s0 = 10
+        mean = 0.1
+        sigma = 0.3
+        dtype = jnp.float32
+        gbm = UnivGeometricBrownianMotion(s0, mean, sigma, dtype)
+
+        stock_paths = gbm.simulate_paths(SEED, 1.0, 52, 100)
+
+        assert stock_paths.shape == (52, 100)
+
 
 class TestMultiGBM:
 
@@ -33,6 +45,7 @@ def test_init(self):
         assert jnp.array_equal(gbm.corr, corr)
         assert gbm.dtype == dtype
         assert jnp.array_equal(gbm.s0, s0)
+        assert gbm.dimension == 2
 
 
     def test_sample_path(self):
@@ -42,6 +55,6 @@ def test_sample_path(self):
         corr = jnp.array([[1, 0.1], [0.1, 1]])
         dtype = jnp.float32
         gbm = MultiGeometricBrownianMotion(s0, mean, sigma, corr, dtype)
-        sample_path = gbm.simulate_paths(0, 1, 52)
+        sample_path = gbm.simulate_paths(SEED, 1.0, 52, 100)
 
-        assert True
+        assert sample_path.shape == (52, 100, 2)

From 0812012b1566adbdb80608e4436f60725c41a24d Mon Sep 17 00:00:00 2001
From: Paolo D'Elia <paolo.delia99@gmail.com>
Date: Tue, 20 Feb 2024 15:12:49 +0100
Subject: [PATCH 5/6] Runned isort

---
 jaxfin/models/gbm/__init__.py                        | 2 +-
 jaxfin/price_engine/black_scholes/__init__.py        | 2 +-
 jaxfin/price_engine/black_scholes/vanilla_options.py | 2 +-
 jaxfin/price_engine/common.py                        | 3 ++-
 4 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/jaxfin/models/gbm/__init__.py b/jaxfin/models/gbm/__init__.py
index 18bff58..786774e 100644
--- a/jaxfin/models/gbm/__init__.py
+++ b/jaxfin/models/gbm/__init__.py
@@ -2,8 +2,8 @@
 Geometric Brownian motion module
 """
 from jaxfin.models.gbm.gbm import (
-    UnivGeometricBrownianMotion,
     MultiGeometricBrownianMotion,
+    UnivGeometricBrownianMotion,
 )
 
 __all__ = ["UnivGeometricBrownianMotion", "MultiGeometricBrownianMotion"]
diff --git a/jaxfin/price_engine/black_scholes/__init__.py b/jaxfin/price_engine/black_scholes/__init__.py
index a6852f4..0590fb3 100644
--- a/jaxfin/price_engine/black_scholes/__init__.py
+++ b/jaxfin/price_engine/black_scholes/__init__.py
@@ -1,7 +1,7 @@
 """Option prices computed with the Black-Scholes module"""
 from jaxfin.price_engine.black_scholes.vanilla_options import (
-    delta_vanilla,
     bs_price,
+    delta_vanilla,
     gamma_vanilla,
 )
 
diff --git a/jaxfin/price_engine/black_scholes/vanilla_options.py b/jaxfin/price_engine/black_scholes/vanilla_options.py
index d8b3c3a..c753687 100644
--- a/jaxfin/price_engine/black_scholes/vanilla_options.py
+++ b/jaxfin/price_engine/black_scholes/vanilla_options.py
@@ -5,7 +5,7 @@
 import jax.numpy as jnp
 
 from ..common import compute_undiscounted_call_prices
-from ..math import d1, cum_normal, density_normal
+from ..math import cum_normal, d1, density_normal
 from ..utils import cast_arrays
 
 
diff --git a/jaxfin/price_engine/common.py b/jaxfin/price_engine/common.py
index 13477f8..d893160 100644
--- a/jaxfin/price_engine/common.py
+++ b/jaxfin/price_engine/common.py
@@ -4,7 +4,8 @@
 
 import jax
 import jax.numpy as jnp
-from .math import d1, cum_normal
+
+from .math import cum_normal, d1
 
 
 def _compute_d1_d2(spots, strikes, expires, vols, discount_rates):

From 899542d678b07166f690b1be2e3581380141a687 Mon Sep 17 00:00:00 2001
From: Paolo D'Elia <paolo.delia99@gmail.com>
Date: Tue, 20 Feb 2024 15:17:45 +0100
Subject: [PATCH 6/6] Runned isort

---
 jaxfin/price_engine/black_scholes/vanilla_options.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/jaxfin/price_engine/black_scholes/vanilla_options.py b/jaxfin/price_engine/black_scholes/vanilla_options.py
index 6c84124..5b2df33 100644
--- a/jaxfin/price_engine/black_scholes/vanilla_options.py
+++ b/jaxfin/price_engine/black_scholes/vanilla_options.py
@@ -6,7 +6,6 @@
 from jax import grad, vmap
 
 from ..common import compute_discounted_call_prices
-
 from ..utils import cast_arrays