fixing typing

dev/lint

@@ -1,6 +1,14 @@
 #!/usr/bin/env bash
+set -e
 
-isort . $1
-black . --exclude "build|dev|.venv|.jupyter|.ipynb_checkpoints" $1
-flake8
-mypy quantflow/utils quantflow/data
+BLACK_ARG="--check"
+RUFF_ARG=""
+
+if [ "$1" = "fix" ] ; then
+    BLACK_ARG=""
+    RUFF_ARG="--fix"
+fi
+
+black quantflow tests ${BLACK_ARG}
+ruff quantflow tests ${RUFF_ARG}
+mypy quantflow

pyproject.toml

@@ -6,21 +6,18 @@ authors = ["Luca <[email protected]>"]
 license = "BSD-3-Clause"
 
 [tool.poetry.dependencies]
-python = "^3.10,<4"
+python = ">=3.10,<3.12"
 numpy = "^1.22.3"
-scipy = "^1.8.0"
-pandas = "^1.4.2"
+scipy = "^1.10.1"
+pandas = "^2.0.1"
 aiohttp = {version = "^3.8.1", optional = true}
 
 [tool.poetry.group.dev.dependencies]
-black = "^22.3.0"
-flake8 = "^6.0.0"
-isort = "^5.10.1"
-flake8-blind-except = "^0.2.1"
-flake8-builtins = "^2.1.0"
+black = "^23.3.0"
 pytest-cov = "^4.0.0"
-mypy = "^0.991"
+mypy = "^1.3.0"
 ghp-import = "^2.0.2"
+ruff = "^0.0.269"
 
 [tool.poetry.extras]
 data = ["aiohttp"]

quantflow/sp/base.py

@@ -1,5 +1,7 @@
+from __future__ import annotations
+
 from abc import ABC, abstractmethod
-from typing import Tuple
+from typing import Generic, Tuple, TypeVar
 
 import numpy as np
 from scipy.optimize import Bounds
@@ -17,7 +19,9 @@ class StochasticProcess(ABC):
     Base class for stochastic processes in continuous time
     """
 
-    parameters: Parameters = Parameters()
+    @property
+    def parameters(self) -> Parameters:
+        return Parameters()
 
     def sample(self, n: int, t: float = 1, steps: int = 0) -> np.ndarray:
         """Generate random paths from the process
@@ -57,26 +61,6 @@ def __str__(self) -> str:
         return self.__repr__()
 
 
-class StochasticProcess1DMarginal(Marginal1D):
-    def __init__(self, process: "StochasticProcess1D", t: float, N: int) -> None:
-        self.process = process
-        self.t = t
-        self.N = N
-
-    def pdf(self, n: Vector) -> Vector:
-        return self.process.pdf(self.t, n)
-
-    def std_norm(self) -> float:
-        """Standard deviation at a time horizon normalized by the time"""
-        return np.sqrt(self.variance() / self.t)
-
-    def characteristic(self, u: Vector) -> Vector:
-        return self.process.characteristic(self.t, u)
-
-    def domain_range(self) -> Bounds:
-        return self.process.domain_range()
-
-
 class StochasticProcess1D(StochasticProcess):
     """
     Base class for 1D stochastic process in continuous time
@@ -177,6 +161,29 @@ def domain_range(self) -> Bounds:
         return default_bounds()
 
 
+P = TypeVar("P", bound=StochasticProcess1D)
+
+
+class StochasticProcess1DMarginal(Marginal1D, Generic[P]):
+    def __init__(self, process: P, t: float, N: int) -> None:
+        self.process = process
+        self.t = t
+        self.N = N
+
+    def pdf(self, n: Vector) -> Vector:
+        return self.process.pdf(self.t, n)
+
+    def std_norm(self) -> float:
+        """Standard deviation at a time horizon normalized by the time"""
+        return np.sqrt(self.variance() / self.t)
+
+    def characteristic(self, u: Vector) -> Vector:
+        return self.process.characteristic(self.t, u)
+
+    def domain_range(self) -> Bounds:
+        return self.process.domain_range()
+
+
 class CountingProcess1D(StochasticProcess1D):
     pass
 
@@ -238,7 +245,7 @@ def difference_process(self) -> CountingProcess1D:
         variables of the process.
         """
 
-    def cdf_square(self, t: float, n: int) -> np.array:
+    def cdf_square(self, t: float, n: int) -> np.ndarray:
         """Cumulative distribution function on a n x n square support"""
         raise NotImplementedError
 

quantflow/sp/ou.py

@@ -59,7 +59,7 @@ def cumulative_characteristic(self, t: float, u: Vector) -> Vector:
         )
         return np.exp(c0 + c1 * self.rate.value)
 
-    def sample(self, n: int, t: float = 1, steps: int = 0) -> np.array:
+    def sample(self, n: int, t: float = 1, steps: int = 0) -> np.ndarray:
         size, dt = self.sample_dt(t, steps)
         jump_process = self.jump_process
         paths = np.zeros((size + 1, n))

quantflow/sp/poisson.py

@@ -1,5 +1,3 @@
-from typing import List, Tuple
-
 import numpy as np
 from scipy.stats import poisson, skellam
 
@@ -22,7 +20,7 @@ def __init__(self, rate: float) -> None:
             "lambda", rate, bounds=(0, None), description="intensity rate"
         )
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return f"{self.__class__.__name__} {self.rate}"
 
     @property
@@ -78,7 +76,7 @@ def characteristic_exponent(self, u: Vector) -> Vector:
     def characteristic(self, t: float, u: Vector) -> Vector:
         return np.exp(t * self.characteristic_exponent(u))
 
-    def sample(self, n: int, t: float = 1, steps: int = 0) -> np.array:
+    def sample(self, n: int, t: float = 1, steps: int = 0) -> np.ndarray:
         size, dt = self.sample_dt(t, steps)
         paths = np.zeros((size + 1, n))
         for p in range(n):
@@ -95,18 +93,22 @@ def sample(self, n: int, t: float = 1, steps: int = 0) -> np.array:
                 paths[i:, p] = y
         return paths
 
-    def arrivals(self, t: float = 1) -> List[float]:
+    def arrivals(self, t: float = 1) -> list[float]:
         """Generate a list of jump arrivals times up to time t"""
         exp_rate = 1.0 / self.rate.value
         arrivals = []
-        tt = 0
+        tt = 0.0
         while tt <= t:
             arrivals.append(tt)
             dt = np.random.exponential(scale=exp_rate)
             tt += dt
         return arrivals
 
-    def jumps(self, n: int) -> np.array:
+    def jumps(self, n: int) -> np.ndarray:
+        """Generate a list of jump sizes
+
+        For a poisson process this is just a list of 1s
+        """
         return np.ones((n,))
 
 
@@ -152,7 +154,7 @@ def pdf(self, t: float, n: Vector = 0) -> Vector:
         """
         return skellam.pmf(n, t * self.rate_left, t * self.rate_right)
 
-    def sample(self, n: int, t: float = 1, steps: int = 0) -> np.array:
+    def sample(self, n: int, t: float = 1, steps: int = 0) -> np.ndarray:
         raise NotImplementedError
 
 
@@ -178,7 +180,7 @@ def __init__(self, rate_left: float, rate_right: float):
         self.rate_left = rate_left
         self.rate_right = rate_right
 
-    def pdf(self, t: float, n: Tuple[Vector, Vector]) -> Vector:
+    def pdf(self, t: float, n: tuple[Vector, Vector]) -> Vector:
         """
         PDF of the process. It's just the product of two
         poisson pdfs :eq:`poisson_pdf`.
@@ -187,7 +189,7 @@ def pdf(self, t: float, n: Tuple[Vector, Vector]) -> Vector:
             n[1], t * self.rate_right
         )
 
-    def cdf(self, t: float, n: Tuple[Vector, Vector]) -> Vector:
+    def cdf(self, t: float, n: tuple[Vector, Vector]) -> Vector:
         """
         CDF of the process. It's just the product of two
         poisson cdfs :eq:`poisson_cdf`.
@@ -196,7 +198,7 @@ def cdf(self, t: float, n: Tuple[Vector, Vector]) -> Vector:
             n[1], t * self.rate_right
         )
 
-    def marginals(self) -> Tuple[CountingProcess1D, CountingProcess1D]:
+    def marginals(self) -> tuple[CountingProcess1D, CountingProcess1D]:
         """
         Returns the marginal poisson processes of each of the two random variables
         """
@@ -216,6 +218,6 @@ def difference_process(self) -> CountingProcess1D:
         """
         return SkellamProcess(self.rate_left, self.rate_right)
 
-    def sample(self, n: int, t: float = 1, steps: int = 0) -> np.array:
+    def sample(self, n: int, t: float = 1, steps: int = 0) -> np.ndarray:
         """require implementation"""
         raise NotImplementedError

quantflow/sp/weibull.py

@@ -1,13 +1,15 @@
 """Weibull model"""
 from functools import lru_cache
 from math import floor
+from typing import cast
 
 import numpy as np
 from numpy import arange, array, exp, power
 from numpy import sum as npsum
 from scipy.special import gamma, gammaln
 
-from ..utils.types import Vector
+from quantflow.utils.types import Vector
+
 from .base import CountingProcess1D
 
 
@@ -68,13 +70,16 @@ def pdf(self, t: float, n: Vector) -> Vector:
         """
         if isinstance(n, np.ndarray):
             return np.array([self.pdf(t, i) for i in n.ravel()]).reshape(n.shape)
-        j = arange(n, self.N)
-        return npsum(
-            power(-1, j + n)
-            * power(self.la * power(t, self.c), j)
-            * self.alpha(n)
-            / gamma(self.c * j + 1)
-        )
+        elif isinstance(n, int):
+            j = arange(n, self.N)
+            return npsum(
+                power(-1, j + n)
+                * power(self.la * power(t, self.c), j)
+                * self.alpha(n)
+                / gamma(self.c * j + 1)
+            )
+        else:
+            raise TypeError("n must be an integer or array of integers")
 
     def cdf(self, t: float, n: Vector) -> Vector:
         r"""
@@ -84,10 +89,13 @@ def cdf(self, t: float, n: Vector) -> Vector:
         """
         if isinstance(n, np.ndarray):
             return np.array([self.cdf(t, i) for i in n.ravel()]).reshape(n.shape)
-        return self.pdf(t, np.arange(0, floor(n) + 1)).sum()
+        elif isinstance(n, int):
+            return cast(np.ndarray, self.pdf(t, np.arange(0, floor(n) + 1))).sum()
+        else:
+            raise TypeError("n must be an integer or array of integers")
 
     @lru_cache(maxsize=None)
-    def alpha(self, n: int) -> array:
+    def alpha(self, n: int) -> np.ndarray:
         """Calculate the alpha coefficients for n events"""
         J = arange(n, self.N)
         if n == 0:

quantflow/sp/weiner.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from functools import lru_cache
 
 import numpy as np
@@ -8,31 +10,6 @@
 from .base import StochasticProcess1D, StochasticProcess1DMarginal
 
 
-class WeinerMarginal(StochasticProcess1DMarginal):
-    def variance(self) -> float:
-        s = self.process.sigma.value
-        return s * s * self.t
-
-    @lru_cache
-    def create_pdf(self) -> Vector:
-        u = -2 * np.pi * np.fft.rfftfreq(self.N)
-        psi = self.characteristic(u)
-        return np.fft.irfft(psi)
-
-    def pdf(self, n: Vector) -> Vector:
-        return norm.pdf(n, scale=self.std())
-
-    def cdf(self, n: Vector) -> Vector:
-        """
-        Compute the cumulative distribution function of the process.
-
-        :param t: time horizon
-        :param n: Location in the stochastic process domain space. If a numpy array,
-            the output should have the same shape as the input.
-        """
-        return n
-
-
 class Weiner(StochasticProcess1D):
     r"""The Heston stochastic volatility model
 
@@ -62,5 +39,27 @@ def characteristic(self, t: float, u: Vector) -> Vector:
         su = self.sigma.value * u
         return np.exp(-0.5 * su * su * t)
 
-    def cdf(self):
-        pass
+
+class WeinerMarginal(StochasticProcess1DMarginal[Weiner]):
+    def variance(self) -> float:
+        s = self.process.sigma.value
+        return s * s * self.t
+
+    @lru_cache
+    def create_pdf(self) -> Vector:
+        u = -2 * np.pi * np.fft.rfftfreq(self.N)
+        psi = self.characteristic(u)
+        return np.fft.irfft(psi)
+
+    def pdf(self, n: Vector) -> Vector:
+        return norm.pdf(n, scale=self.std())
+
+    def cdf(self, n: Vector) -> Vector:
+        """
+        Compute the cumulative distribution function of the process.
+
+        :param t: time horizon
+        :param n: Location in the stochastic process domain space. If a numpy array,
+            the output should have the same shape as the input.
+        """
+        return n