chore: Improve type suggestions

README.md

@@ -1,3 +1,29 @@
 # Affine Diffusions
 
 Simulation and estimation of Affine Diffusion models.
+
+Install:
+
+```shell
+pip install affidiff
+```
+
+## Contribute
+
+Install project in editable mode and sync all dependencies:
+
+```shell
+uv sync --all-groups
+```
+
+and use pre-commit to make sure that your code is formatted and linted automatically:
+
+```shell
+uv run prek install
+```
+
+Run tests:
+
+```shell
+uv run pytest
+```

examples/try_centtend.py

@@ -1,9 +1,7 @@
-# ruff: noqa
 """Try Central Tendency model."""
 
 from __future__ import annotations
 
-
 import itertools
 import time
 
@@ -74,7 +72,7 @@ def try_simulation_pq() -> None:
  eta_y=eta_y,
  rho=rho,
  )
-centtend = CentTend(param_true)
+_ = CentTend(param_true)
  print(param_true)
  print(param_true.is_valid())
 
@@ -164,11 +162,8 @@ def try_sim_realized() -> None:
  plt.show()
 
 
-def try_sim_realized_pq():
- """Simulate realized data from Central Tendency model
- under P and Q measures.
-
- """
+def try_sim_realized_pq() -> None:
+ """Simulate realized data from Central Tendency model under P and Q measures."""
  riskfree = 0.0
  lmbd = 1.01
  lmbd_s, lmbd_y = 0.5, 0.5
@@ -209,7 +204,7 @@ def try_sim_realized_pq():
  )
 
 
-def try_integrated_gmm_single():
+def try_integrated_gmm_single() -> None:
  """Simulate realized data from Central Tendency model. Estimate parameters."""
  riskfree = 0.0
  mean_v = 0.2
@@ -274,7 +269,7 @@ def try_integrated_gmm_single():
  print("Elapsed time = %.2f min" % ((time.time() - time_start) / 60))
 
 
-def try_integrated_gmm_real():
+def try_integrated_gmm_real() -> None:
  """Estimate Central Tendency model parameters with real data."""
  riskfree = 0.0
 
@@ -327,8 +322,9 @@ def try_integrated_gmm_real():
  print("Elapsed time = %.2f min" % ((time.time() - time_start) / 60))
 
 
-def try_integrated_gmm_opt_methods():
+def try_integrated_gmm_opt_methods() -> None:
  """Simulate realized data from Central Tendency model. Estimate parameters.
+
  Check various optimization methods.
 
  """

examples/try_heston.py

@@ -1,4 +1,3 @@
-# ruff: noqa
 """Try Heston model."""
 
 from __future__ import annotations
@@ -10,13 +9,14 @@
 import numpy as np
 import seaborn as sns
 from load_real_data import load_data # type: ignore
+from mygmm import Results
 from statsmodels.tsa.stattools import acf
 
 from affidiff import Heston, HestonParam
 from affidiff.helper_functions import plot_final_distr, plot_realized, plot_trajectories, take_time
 
 
-def try_simulation():
+def try_simulation() -> None:
  """Try simulating and plotting Heston model."""
  riskfree = 0.0
  lmbd = 0.0
@@ -40,7 +40,7 @@ def try_simulation():
  plot_trajectories(paths=volatility, nsub=nsub, names="volatility")
 
 
-def try_simulation_pq():
+def try_simulation_pq() -> None:
  """Try simulating and plotting Heston model."""
  riskfree = 0.0
  lmbd = 0.0
@@ -76,7 +76,7 @@ def try_simulation_pq():
  plot_trajectories(paths=[volatility, volatility_q], nsub=nsub, names=["volatility", "volatility_q"])
 
 
-def try_marginal():
+def try_marginal() -> None:
  """Simulate and plot marginal distribution of the data in Heston model."""
  riskfree = 0.0
  lmbd = 0.0
@@ -100,7 +100,7 @@ def try_marginal():
  plot_final_distr(paths=volatility, names="volatility")
 
 
-def try_sim_realized():
+def try_sim_realized() -> None:
  """Simulate realized data from Heston model and plot it."""
  riskfree = 0.0
  lmbd = 0.0
@@ -121,7 +121,7 @@ def try_sim_realized():
  plot_realized(returns=returns, rvar=rvar)
 
 
-def try_sim_realized_pq():
+def try_sim_realized_pq() -> None:
  """Simulate realized data from Heston model under P and Q measures."""
  riskfree = 0.0
  mean_v = 0.5
@@ -151,7 +151,7 @@ def try_sim_realized_pq():
  )
 
 
-def try_integrated_gmm_single():
+def try_integrated_gmm_single() -> None:
  """Simulate realized data from Heston model. Estimate parameters."""
  riskfree = 0.0
 
@@ -196,8 +196,9 @@ def try_integrated_gmm_single():
  print("Elapsed time = %.2f min" % ((time.time() - time_start) / 60))
 
 
-def try_integrated_gmm_single_rn():
+def try_integrated_gmm_single_rn() -> None:
  """Simulate realized data from risk-neutral Heston model.
+
  Estimate parameters.
 
  """
@@ -261,8 +262,9 @@ def try_integrated_gmm_single_rn():
  print("Elapsed time = %.2f min" % ((time.time() - time_start) / 60))
 
 
-def try_integrated_gmm_joint():
+def try_integrated_gmm_joint() -> Results:
  """Simulate realized data from risk-neutral Heston model.
+
  Estimate parameters.
 
  """
@@ -313,7 +315,7 @@ def try_integrated_gmm_joint():
  return res
 
 
-def try_integrated_gmm_real():
+def try_integrated_gmm_real() -> None:
  """Estimate Heston model parameters with real data."""
  riskfree = 0.0
 
@@ -353,8 +355,9 @@ def try_integrated_gmm_real():
  print("Elapsed time = %.2f min" % ((time.time() - time_start) / 60))
 
 
-def try_integrated_gmm_opt_methods():
+def try_integrated_gmm_opt_methods() -> None:
  """Simulate realized data from Heston model. Estimate parameters.
+
  Check various optimization methods.
 
  """

src/affidiff/helper_functions.py

@@ -5,15 +5,18 @@
 import contextlib
 import itertools as it
 import time
-from typing import Any, Callable, Generator, Sequence
+from typing import TYPE_CHECKING, Any, Callable, Generator, Sequence
+
+if TYPE_CHECKING:
+ from affidiff.param_generic import GenericParam
 
 import matplotlib.pylab as plt
 import numpy as np
 import seaborn as sns
 from statsmodels.tsa.tsatools import lagmat
 
 
-def ajd_drift(*, state: Any, theta: Any) -> np.ndarray: # noqa: ANN401
+def ajd_drift(*, state: np.ndarray | float | Sequence[float], theta: GenericParam) -> np.ndarray:
  """Instantaneous mean.
 
  Parameters
@@ -33,7 +36,7 @@ def ajd_drift(*, state: Any, theta: Any) -> np.ndarray: # noqa: ANN401
  return theta.mat_k0 + state_arr.dot(np.transpose(theta.mat_k1))
 
 
-def ajd_diff(*, state: Any, theta: Any) -> np.ndarray: # noqa: ANN401
+def ajd_diff(*, state: np.ndarray | float | Sequence[float], theta: GenericParam) -> np.ndarray:
  """Instantaneous volatility.
 
  Parameters
@@ -81,7 +84,7 @@ def nice_errors(*, errors: np.ndarray, sdim: int) -> np.ndarray:
  return np.concatenate((errors, -errors), axis=sdim)
 
 
-def plot_trajectories(*, paths: Any, nsub: int, names: str | list[str]) -> None: # noqa: ANN401
+def plot_trajectories(*, paths: np.ndarray | list[np.ndarray], nsub: int, names: str | list[str]) -> None:
  """Plot process realizations.
 
  Parameters
@@ -111,7 +114,7 @@ def plot_trajectories(*, paths: Any, nsub: int, names: str | list[str]) -> None:
  plt.show()
 
 
-def plot_final_distr(*, paths: Any, names: str | list[str]) -> None: # noqa: ANN401
+def plot_final_distr(*, paths: np.ndarray | list[np.ndarray], names: str | list[str]) -> None:
  """Plot marginal distribution of the process.
 
  Parameters
@@ -140,8 +143,8 @@ def plot_final_distr(*, paths: Any, names: str | list[str]) -> None: # noqa: AN
 
 def plot_realized(
  *,
- returns: Any, # noqa: ANN401
- rvar: Any, # noqa: ANN401
+ returns: np.ndarray | list[np.ndarray],
+ rvar: np.ndarray | list[np.ndarray],
  suffix: list[str] | None = None,
 ) -> None:
  """Plot realized returns and volatility.
@@ -287,7 +290,7 @@ def poly_coef(roots: Sequence[float] | np.ndarray) -> list[float]:
 
 def instruments(
  *,
- data: Any = None, # noqa: ANN401
+ data: np.ndarray | None = None,
  instrlag: int = 1,
  nobs: int | None = None,
  instr_choice: str = "const",

src/affidiff/model_cir.py

@@ -2,18 +2,22 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING
 
 from affidiff.model_generic import SDE
 
 if TYPE_CHECKING:
  import numpy as np
 
+ from affidiff.param_cir import CIRparam
+
 
 class CIR(SDE):
  """Cox-Ingersoll-Ross (CIR) model."""
 
- def __init__(self, param: Any = None) -> None: # noqa: ANN401
+ param: CIRparam | None
+
+ def __init__(self, param: CIRparam | None = None) -> None:
  """Initialize the class.
 
  Parameters
@@ -33,10 +37,11 @@ def get_start(self) -> list[float]:
  Starting value at the long-run mean
 
  """
+ assert self.param is not None
  return [float(self.param.mean)]
 
  @staticmethod
- def drift(*, state: np.ndarray | float, theta: Any) -> np.ndarray | float: # noqa: ANN401
+ def drift(*, state: np.ndarray | float, theta: CIRparam) -> np.ndarray | float:
  """Drift function.
 
  Parameters
@@ -55,7 +60,7 @@ def drift(*, state: np.ndarray | float, theta: Any) -> np.ndarray | float: # no
  return theta.kappa * (theta.mean - state)
 
  @staticmethod
- def diff(*, state: np.ndarray | float, theta: Any) -> np.ndarray | float: # noqa: ANN401
+ def diff(*, state: np.ndarray | float, theta: CIRparam) -> np.ndarray | float:
  """Diffusion (instantaneous volatility) function.
 
  Parameters