From b190a60c42950249403a7bb80c8e20d269cad9b6 Mon Sep 17 00:00:00 2001
From: Stanislav Khrapov <khrapovs@gmail.com>
Date: Thu, 30 Jul 2026 20:37:02 +0200
Subject: [PATCH 1/7] docs: add installation, contribution, and testing
 instructions to README

---
 README.md | 26 ++++++++++++++++++++++++++
 1 file changed, 26 insertions(+)

diff --git a/README.md b/README.md
index ccbdb7b..edd9a42 100644
--- a/README.md
+++ b/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
+```

From e79086597863f12dbbab689f28c81e60975d26e6 Mon Sep 17 00:00:00 2001
From: Stanislav Khrapov <khrapovs@gmail.com>
Date: Thu, 30 Jul 2026 20:53:13 +0200
Subject: [PATCH 2/7] refactor: improve type hinting and update docstrings
 across model modules

---
 src/affidiff/model_gbm.py     |  2 +-
 src/affidiff/model_generic.py | 36 +++++++++++++++++------------------
 2 files changed, 19 insertions(+), 19 deletions(-)

diff --git a/src/affidiff/model_gbm.py b/src/affidiff/model_gbm.py
index 7621d17..8f443fc 100644
--- a/src/affidiff/model_gbm.py
+++ b/src/affidiff/model_gbm.py
@@ -297,7 +297,7 @@ def integrated_mom(
             Model parameters
         data : (2, nobs) array
             Returns and realized variance
-        instr_data : object
+        instr_data : array_like, optional
             Instrument data
         instr_choice : str
             Instrument choice
diff --git a/src/affidiff/model_generic.py b/src/affidiff/model_generic.py
index 294ce7f..04c2c91 100644
--- a/src/affidiff/model_generic.py
+++ b/src/affidiff/model_generic.py
@@ -57,7 +57,7 @@ def __init__(self, param: Any = None) -> None:  # noqa: ANN401
         self.param: Any = param
         self.errors: np.ndarray | None = None
 
-    def update_theta(self, param: GenericParam | object) -> None:
+    def update_theta(self, param: GenericParam) -> None:
         """Update model parameters.
 
         Parameters
@@ -92,7 +92,7 @@ def realized_const(
         """Realized constant in integrated moments."""
         raise NotImplementedError("Must be overridden")
 
-    def euler_loc(self, *, state: np.ndarray, theta: GenericParam | object) -> np.ndarray:
+    def euler_loc(self, *, state: np.ndarray, theta: GenericParam) -> np.ndarray:
         """Euler location.
 
         Parameters
@@ -110,7 +110,7 @@ def euler_loc(self, *, state: np.ndarray, theta: GenericParam | object) -> np.nd
         """
         return ajd_drift(state=state, theta=theta)
 
-    def euler_scale(self, *, state: np.ndarray, theta: GenericParam | object) -> np.ndarray:
+    def euler_scale(self, *, state: np.ndarray, theta: GenericParam) -> np.ndarray:
         """Euler scale.
 
         Parameters
@@ -128,7 +128,7 @@ def euler_scale(self, *, state: np.ndarray, theta: GenericParam | object) -> np.
         """
         return ajd_diff(state=state, theta=theta)
 
-    def loc(self, *, state: np.ndarray, theta: GenericParam | object) -> np.ndarray:
+    def loc(self, *, state: np.ndarray, theta: GenericParam) -> np.ndarray:
         """Location.
 
         Parameters
@@ -146,7 +146,7 @@ def loc(self, *, state: np.ndarray, theta: GenericParam | object) -> np.ndarray:
         """
         return self.euler_loc(state=state, theta=theta)
 
-    def scale(self, *, state: np.ndarray, theta: GenericParam | object) -> np.ndarray:
+    def scale(self, *, state: np.ndarray, theta: GenericParam) -> np.ndarray:
         """Scale.
 
         Parameters
@@ -164,11 +164,11 @@ def scale(self, *, state: np.ndarray, theta: GenericParam | object) -> np.ndarra
         """
         return self.euler_scale(state=state, theta=theta)
 
-    def exact_loc(self, *, state: np.ndarray, theta: GenericParam | object) -> np.ndarray:
+    def exact_loc(self, *, state: np.ndarray, theta: GenericParam) -> np.ndarray:
         """Exact location."""
         return self.euler_loc(state=state, theta=theta)
 
-    def exact_scale(self, *, state: np.ndarray, theta: GenericParam | object) -> np.ndarray:
+    def exact_scale(self, *, state: np.ndarray, theta: GenericParam) -> np.ndarray:
         """Exact scale."""
         return self.euler_scale(state=state, theta=theta)
 
@@ -202,7 +202,7 @@ def momcond(
         """Moment conditions."""
         raise NotImplementedError
 
-    def depvar_unc_mean(self, *, param: GenericParam | object, aggh: float) -> np.ndarray:
+    def depvar_unc_mean(self, *, param: GenericParam, aggh: float) -> np.ndarray:
         """Unconditional means of realized data.
 
         Parameters
@@ -484,7 +484,7 @@ def gmmest(
         method: str = "BFGS",
         kernel: str = "Bartlett",
         band: int | None = None,
-    ) -> object:
+    ) -> Any:  # noqa: ANN401
         """Estimate model parameters using GMM.
 
         Parameters
@@ -523,22 +523,22 @@ def gmmest(
     def integrated_gmm(
         self,
         *,
-        param_start: GenericParam | object,
-        data: object = None,
-        instr_data: object = None,
+        param_start: GenericParam,
+        data: Any = None,  # noqa: ANN401
+        instr_data: Any = None,  # noqa: ANN401
         instr_choice: str = "const",
-        aggh: object = 1,
+        aggh: float | Sequence[float] = 1,
         instrlag: int = 1,
         subset: str = "all",
         measure: str = "P",
         names: list[str] | None = None,
         bounds: list[tuple[float | None, float | None]] | None = None,
-        constraints: object = (),
+        constraints: Any = (),  # noqa: ANN401
         iter: int = 2,
         method: str = "BFGS",
         kernel: str = "Bartlett",
         band: int | None = None,
-    ) -> object:
+    ) -> Any:  # noqa: ANN401
         """Estimate model parameters using Integrated GMM.
 
         Parameters
@@ -615,10 +615,10 @@ def integrated_mom(
         self,
         *,
         theta: np.ndarray | Sequence[float],
-        data: object = None,
-        instr_data: object = None,
+        data: Any = None,  # noqa: ANN401
+        instr_data: Any = None,  # noqa: ANN401
         instr_choice: str = "const",
-        aggh: object = 1,
+        aggh: float | Sequence[float] = 1,
         subset: str = "all",
         instrlag: int = 1,
         measure: str = "P",

From 488ba4d5a54948f81a506307ea763a312e099e42 Mon Sep 17 00:00:00 2001
From: Stanislav Khrapov <khrapovs@gmail.com>
Date: Thu, 30 Jul 2026 21:02:29 +0200
Subject: [PATCH 3/7] refactor: improve type safety and add explicit parameter
 typing across all SDE models and helper functions

---
 src/affidiff/helper_functions.py | 19 ++++---
 src/affidiff/model_cir.py        | 13 +++--
 src/affidiff/model_ct.py         | 86 +++++++++++++++++---------------
 src/affidiff/model_gbm.py        | 64 +++++++++++++++---------
 src/affidiff/model_generic.py    | 79 +++++++++++++++--------------
 src/affidiff/model_heston.py     | 41 +++++++++------
 src/affidiff/model_vasicek.py    | 13 +++--
 src/affidiff/param_generic.py    |  8 ++-
 8 files changed, 187 insertions(+), 136 deletions(-)

diff --git a/src/affidiff/helper_functions.py b/src/affidiff/helper_functions.py
index dfd0688..4ae079b 100644
--- a/src/affidiff/helper_functions.py
+++ b/src/affidiff/helper_functions.py
@@ -5,7 +5,10 @@
 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
@@ -13,7 +16,7 @@
 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",
diff --git a/src/affidiff/model_cir.py b/src/affidiff/model_cir.py
index dea6652..a56a9c9 100644
--- a/src/affidiff/model_cir.py
+++ b/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
diff --git a/src/affidiff/model_ct.py b/src/affidiff/model_ct.py
index 50181b2..842c262 100644
--- a/src/affidiff/model_ct.py
+++ b/src/affidiff/model_ct.py
@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 from math import exp
-from typing import TYPE_CHECKING, Any, cast
+from typing import TYPE_CHECKING, Sequence, cast
 
 import numpy as np
 from statsmodels.tsa.tsatools import lagmat
@@ -13,15 +13,15 @@
 from affidiff.param_ct import CentTendParam
 
 if TYPE_CHECKING:
-    pass
+    from affidiff.param_generic import GenericParam
 
 
 class CentTend(SDE):
     """Central Tendency model."""
 
-    param: Any
+    param: CentTendParam | None
 
-    def __init__(self, param: Any = None) -> None:  # noqa: ANN401
+    def __init__(self, param: CentTendParam | None = None) -> None:
         """Initialize the class.
 
         Parameters
@@ -35,18 +35,12 @@ def __init__(self, param: Any = None) -> None:  # noqa: ANN401
         super().__init__(param)
 
     def get_start(self) -> list[float]:
-        """Get starting values for simulation.
-
-        Returns
-        -------
-        array_like
-            Starting values for price and variance
-
-        """
+        """Return starting values for simulation."""
+        assert self.param is not None
         return [1.0, float(self.param.mean_v), float(self.param.mean_v)]
 
     @staticmethod
-    def coef_big_as(*, param: Any, aggh: float) -> float:  # noqa: ANN401
+    def coef_big_as(*, param: CentTendParam, aggh: float) -> float:
         r"""Coefficient A^\sigma_h in exact discretization of volatility.
 
         Parameters
@@ -64,7 +58,7 @@ def coef_big_as(*, param: Any, aggh: float) -> float:  # noqa: ANN401
         """
         return float(np.exp(-param.kappa_s * aggh))
 
-    def coef_big_bs(self, *, param: Any, aggh: float) -> float:  # noqa: ANN401
+    def coef_big_bs(self, *, param: CentTendParam, aggh: float) -> float:
         r"""Coefficient B^\sigma_h in exact discretization of volatility.
 
         Parameters
@@ -87,7 +81,7 @@ def coef_big_bs(self, *, param: Any, aggh: float) -> float:  # noqa: ANN401
             * (self.coef_big_ay(param=param, aggh=aggh) - self.coef_big_as(param=param, aggh=aggh))
         )
 
-    def coef_big_cs(self, *, param: Any, aggh: float) -> float:  # noqa: ANN401
+    def coef_big_cs(self, *, param: CentTendParam, aggh: float) -> float:
         """Coefficient C^s_h in exact discretization of volatility.
 
         Parameters
@@ -108,7 +102,7 @@ def coef_big_cs(self, *, param: Any, aggh: float) -> float:  # noqa: ANN401
         )
 
     @staticmethod
-    def coef_big_ay(*, param: Any, aggh: float) -> float:  # noqa: ANN401
+    def coef_big_ay(*, param: CentTendParam, aggh: float) -> float:
         """Coefficient A^v_h in exact discretization of volatility.
 
         Parameters
@@ -125,7 +119,7 @@ def coef_big_ay(*, param: Any, aggh: float) -> float:  # noqa: ANN401
         """
         return float(np.exp(-param.kappa_y * aggh))
 
-    def coef_big_cy(self, *, param: Any, aggh: float) -> float:  # noqa: ANN401
+    def coef_big_cy(self, *, param: CentTendParam, aggh: float) -> float:
         """Coefficient C^y_h in exact discretization of volatility.
 
         Parameters
@@ -143,7 +137,7 @@ def coef_big_cy(self, *, param: Any, aggh: float) -> float:  # noqa: ANN401
         """
         return float(param.mean_v * (1 - self.coef_big_ay(param=param, aggh=aggh)))
 
-    def coef_small_as(self, *, param: Any, aggh: float) -> float:  # noqa: ANN401
+    def coef_small_as(self, *, param: CentTendParam, aggh: float) -> float:
         """Coefficient a^s_h in exact discretization of volatility.
 
         Parameters
@@ -161,7 +155,7 @@ def coef_small_as(self, *, param: Any, aggh: float) -> float:  # noqa: ANN401
         """
         return float((1 - self.coef_big_as(param=param, aggh=aggh)) / param.kappa_s / aggh)
 
-    def coef_small_bs(self, *, param: Any, aggh: float) -> float:  # noqa: ANN401
+    def coef_small_bs(self, *, param: CentTendParam, aggh: float) -> float:
         """Coefficient b^s_h in exact discretization of volatility.
 
         Parameters
@@ -184,7 +178,7 @@ def coef_small_bs(self, *, param: Any, aggh: float) -> float:  # noqa: ANN401
             * (self.coef_small_ay(param=param, aggh=aggh) - self.coef_small_as(param=param, aggh=aggh))
         )
 
-    def coef_small_cs(self, *, param: Any, aggh: float) -> float:  # noqa: ANN401
+    def coef_small_cs(self, *, param: CentTendParam, aggh: float) -> float:
         """Coefficient c^s_h in exact discretization of volatility.
 
         Parameters
@@ -204,7 +198,7 @@ def coef_small_cs(self, *, param: Any, aggh: float) -> float:  # noqa: ANN401
             param.mean_v * (1 - self.coef_small_as(param=param, aggh=aggh) - self.coef_small_bs(param=param, aggh=aggh))
         )
 
-    def coef_small_ay(self, *, param: Any, aggh: float) -> float:  # noqa: ANN401
+    def coef_small_ay(self, *, param: CentTendParam, aggh: float) -> float:
         """Coefficient a^v_h in exact discretization of volatility.
 
         Parameters
@@ -221,7 +215,7 @@ def coef_small_ay(self, *, param: Any, aggh: float) -> float:  # noqa: ANN401
         """
         return float((1 - self.coef_big_ay(param=param, aggh=aggh)) / param.kappa_y / aggh)
 
-    def roots(self, *, param: Any, aggh: float) -> list[float]:  # noqa: ANN401
+    def roots(self, *, param: CentTendParam, aggh: float) -> list[float]:
         r"""Roots of the polynomial in moment restrictions.
 
         .. math::
@@ -253,7 +247,7 @@ def roots(self, *, param: Any, aggh: float) -> list[float]:  # noqa: ANN401
         ]
 
     @staticmethod
-    def mean_vol(*, param: Any, aggh: float) -> float:  # noqa: ARG004, ANN401
+    def mean_vol(*, param: GenericParam, aggh: float) -> float:  # noqa: ARG004
         """Unconditional mean of realized volatiliy.
 
         Parameters
@@ -268,9 +262,10 @@ def mean_vol(*, param: Any, aggh: float) -> float:  # noqa: ARG004, ANN401
         float
 
         """
+        assert isinstance(param, CentTendParam)
         return float(param.mean_v)
 
-    def mean_vol2(self, *, param: Any, aggh: float) -> float:  # noqa: ANN401
+    def mean_vol2(self, *, param: GenericParam, aggh: float) -> float:
         """Unconditional mean of squared realized volatiliy.
 
         Parameters
@@ -285,6 +280,7 @@ def mean_vol2(self, *, param: Any, aggh: float) -> float:  # noqa: ANN401
         float
 
         """
+        assert isinstance(param, CentTendParam)
         return float(
             self.coef_small_as(param=param, aggh=aggh) ** 2 * unc_var_sigma(param)
             + self.coef_small_bs(param=param, aggh=aggh) ** 2 * unc_var_ct(param)
@@ -292,7 +288,7 @@ def mean_vol2(self, *, param: Any, aggh: float) -> float:  # noqa: ANN401
         )
 
     @staticmethod
-    def mean_ret(*, param: Any, aggh: float) -> float:  # noqa: ARG004, ANN401
+    def mean_ret(*, param: GenericParam, aggh: float) -> float:  # noqa: ARG004
         """Unconditional mean of realized returns.
 
         Parameters
@@ -307,9 +303,10 @@ def mean_ret(*, param: Any, aggh: float) -> float:  # noqa: ARG004, ANN401
         float
 
         """
+        assert isinstance(param, CentTendParam)
         return float((param.lmbd - 0.5) * param.mean_v)
 
-    def mean_cross(self, *, param: Any, aggh: float) -> float:  # noqa: ANN401
+    def mean_cross(self, *, param: GenericParam, aggh: float) -> float:
         """Unconditional mean of realized returns times volatility.
 
         Parameters
@@ -324,13 +321,16 @@ def mean_cross(self, *, param: Any, aggh: float) -> float:  # noqa: ANN401
         float
 
         """
+        assert isinstance(param, CentTendParam)
         p = param
         return float(
             (p.lmbd - 0.5) * self.mean_vol2(param=param, aggh=aggh)
             + p.rho * p.mean_v * p.eta_s / p.kappa_s * (1 - self.coef_small_as(param=param, aggh=aggh)) / aggh
         )
 
-    def realized_const(self, *, param: Any = None, aggh: float = 1, subset: slice | None = None) -> np.ndarray:  # noqa: ANN401
+    def realized_const(
+        self, *, param: GenericParam | None = None, aggh: float = 1, subset: slice | None = None
+    ) -> np.ndarray:
         """Intercept in the realized moment conditions.
 
         Parameters
@@ -348,6 +348,9 @@ def realized_const(self, *, param: Any = None, aggh: float = 1, subset: slice |
             Intercept
 
         """
+        if param is None:
+            param = self.param
+        assert isinstance(param, CentTendParam)
         res = (
             (
                 self.mat_a0(param=param, aggh=1)
@@ -363,7 +366,7 @@ def realized_const(self, *, param: Any = None, aggh: float = 1, subset: slice |
             res = res[subset]
         return np.squeeze(res)
 
-    def mat_a0(self, *, param: Any, aggh: float) -> np.ndarray:  # noqa: ANN401
+    def mat_a0(self, *, param: CentTendParam, aggh: float) -> np.ndarray:
         """Matrix A_0 in integrated moments.
 
         Parameters
@@ -383,7 +386,7 @@ def mat_a0(self, *, param: Any, aggh: float) -> np.ndarray:  # noqa: ANN401
         mat[1, 1] = poly_coef(self.roots(param=param, aggh=aggh))[0]
         return mat
 
-    def mat_a1(self, *, param: Any, aggh: float) -> np.ndarray:  # noqa: ANN401
+    def mat_a1(self, *, param: CentTendParam, aggh: float) -> np.ndarray:
         """Matrix A_1 in integrated moments.
 
         Parameters
@@ -403,7 +406,7 @@ def mat_a1(self, *, param: Any, aggh: float) -> np.ndarray:  # noqa: ANN401
         mat[1, 1] = poly_coef(self.roots(param=param, aggh=aggh))[1]
         return mat
 
-    def mat_a2(self, *, param: Any, aggh: float) -> np.ndarray:  # noqa: ANN401
+    def mat_a2(self, *, param: CentTendParam, aggh: float) -> np.ndarray:
         """Matrix A_2 in integrated moments.
 
         Parameters
@@ -423,7 +426,7 @@ def mat_a2(self, *, param: Any, aggh: float) -> np.ndarray:  # noqa: ANN401
         mat[1, 1] = poly_coef(self.roots(param=param, aggh=aggh))[2]
         return mat
 
-    def mat_a3(self, *, param: Any, aggh: float) -> np.ndarray:  # noqa: ANN401
+    def mat_a3(self, *, param: CentTendParam, aggh: float) -> np.ndarray:
         """Matrix A_3 in integrated moments.
 
         Parameters
@@ -446,7 +449,7 @@ def mat_a3(self, *, param: Any, aggh: float) -> np.ndarray:  # noqa: ANN401
         mat[3, 3] = mat[0, 0]
         return mat
 
-    def mat_a4(self, *, param: Any, aggh: float) -> np.ndarray:  # noqa: ANN401
+    def mat_a4(self, *, param: CentTendParam, aggh: float) -> np.ndarray:
         """Matrix A_4 in integrated moments.
 
         Parameters
@@ -469,7 +472,7 @@ def mat_a4(self, *, param: Any, aggh: float) -> np.ndarray:  # noqa: ANN401
         mat[3, 3] = mat[0, 0]
         return mat
 
-    def mat_a5(self, *, param: Any, aggh: float) -> np.ndarray:  # noqa: ANN401
+    def mat_a5(self, *, param: CentTendParam, aggh: float) -> np.ndarray:
         """Matrix A_5 in integrated moments.
 
         Parameters
@@ -494,7 +497,7 @@ def mat_a5(self, *, param: Any, aggh: float) -> np.ndarray:  # noqa: ANN401
         mat[3, 1] = (0.5 - param.lmbd) * mat[0, 0]
         return mat
 
-    def mat_a(self, *, param: Any, subset: slice | None = None) -> np.ndarray:  # noqa: ANN401
+    def mat_a(self, *, param: GenericParam, subset: slice | None = None) -> np.ndarray:
         """Matrix A in integrated moments.
 
         Parameters
@@ -506,10 +509,11 @@ def mat_a(self, *, param: Any, subset: slice | None = None) -> np.ndarray:  # no
 
         Returns
         -------
-        (nmoms, 6*nmoms) array
+        (nmoms, 3*nmoms) array
             Matrix A
 
         """
+        assert isinstance(param, CentTendParam)
         mat_a_tuple = (
             self.mat_a0(param=param, aggh=1),
             self.mat_a1(param=param, aggh=1),
@@ -524,7 +528,7 @@ def mat_a(self, *, param: Any, subset: slice | None = None) -> np.ndarray:  # no
         return np.squeeze(res)
 
     @staticmethod
-    def realized_depvar(*, data: Any, subset: slice | None = None) -> np.ndarray:  # noqa: ANN401
+    def realized_depvar(*, data: np.ndarray | Sequence[np.ndarray], subset: slice | None = None) -> np.ndarray:
         """Array of the left-hand side variables in realized moment conditions.
 
         Parameters
@@ -549,7 +553,7 @@ def realized_depvar(*, data: Any, subset: slice | None = None) -> np.ndarray:  #
         return cast(np.ndarray, lagmat(var_s.T, maxlag=5, original="in"))
 
 
-def unc_mean_ct2(param: Any) -> float:  # noqa: ANN401
+def unc_mean_ct2(param: CentTendParam) -> float:
     """Calculate unconditional second moment of CT, E[y_t**4].
 
     Parameters
@@ -566,7 +570,7 @@ def unc_mean_ct2(param: Any) -> float:  # noqa: ANN401
     return float(p.mean_v * p.eta_y**2 / p.kappa_y / 2)
 
 
-def unc_mean_sigma2(param: Any) -> float:  # noqa: ANN401
+def unc_mean_sigma2(param: CentTendParam) -> float:
     r"""Calculate unconditional second moment of volatility, E[\sigma_t**4].
 
     Parameters
@@ -583,7 +587,7 @@ def unc_mean_sigma2(param: Any) -> float:  # noqa: ANN401
     return float(unc_mean_ct2(param) * p.kappa_s / (p.kappa_s + p.kappa_y) + p.mean_v * p.eta_s**2 / p.kappa_s / 2)
 
 
-def unc_var_ct(param: Any) -> float:  # noqa: ANN401
+def unc_var_ct(param: CentTendParam) -> float:
     """Calculate unconditional variance of CT, V[y_t**2].
 
     Parameters
@@ -600,7 +604,7 @@ def unc_var_ct(param: Any) -> float:  # noqa: ANN401
     return float(p.mean_v**2 + unc_mean_ct2(param))
 
 
-def unc_var_sigma(param: Any) -> float:  # noqa: ANN401
+def unc_var_sigma(param: CentTendParam) -> float:
     r"""Calculate unconditional variance of volatility, V[\sigma_t**2].
 
     Parameters
@@ -617,7 +621,7 @@ def unc_var_sigma(param: Any) -> float:  # noqa: ANN401
     return float(p.mean_v**2 + unc_mean_sigma2(param))
 
 
-def unc_var_error(*, param: Any, aggh: float) -> float:  # noqa: ANN401
+def unc_var_error(*, param: CentTendParam, aggh: float) -> float:
     r"""Calculate unconditional variance of aggregated volatility error.
 
     :math:`V\left[\frac{1}{H}\int_{0}^{H}\epsilon_{t,s}^{\sigma}ds\right]`.
diff --git a/src/affidiff/model_gbm.py b/src/affidiff/model_gbm.py
index 8f443fc..6a862c9 100644
--- a/src/affidiff/model_gbm.py
+++ b/src/affidiff/model_gbm.py
@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Sequence, cast
+from typing import TYPE_CHECKING, Sequence, cast
 
 import numdifftools as nd
 import numpy as np
@@ -13,13 +13,13 @@
 from affidiff.param_gbm import GBMparam
 
 if TYPE_CHECKING:
-    pass
+    from affidiff.param_generic import GenericParam
 
 
 class GBM(SDE):
     """Geometric Brownian Motion."""
 
-    def __init__(self, param: Any = None) -> None:  # noqa: ANN401
+    def __init__(self, param: GBMparam | None = None) -> None:
         """Initialize the class.
 
         Parameters
@@ -35,7 +35,7 @@ def get_start(self) -> list[float]:
         return [1.0]
 
     @staticmethod
-    def drift(*, state: np.ndarray | float, theta: Any) -> np.ndarray | float:  # noqa: ARG004, ANN401
+    def drift(*, state: np.ndarray | float, theta: GBMparam | np.ndarray | Sequence[float]) -> np.ndarray | float:  # noqa: ARG004
         """Drift function.
 
         Parameters
@@ -51,10 +51,12 @@ def drift(*, state: np.ndarray | float, theta: Any) -> np.ndarray | float:  # no
             Drift value
 
         """
-        return theta.mean - theta.sigma**2 / 2
+        if isinstance(theta, GBMparam):
+            return theta.mean - theta.sigma**2 / 2
+        return float(theta[0]) - float(theta[1]) ** 2 / 2
 
     @staticmethod
-    def diff(*, state: np.ndarray | float, theta: Any) -> np.ndarray | float:  # noqa: ARG004, ANN401
+    def diff(*, state: np.ndarray | float, theta: GBMparam | np.ndarray | Sequence[float]) -> np.ndarray | float:  # noqa: ARG004
         """Diffusion (instantaneous volatility) function.
 
         Parameters
@@ -70,7 +72,9 @@ def diff(*, state: np.ndarray | float, theta: Any) -> np.ndarray | float:  # noq
             Diffusion value
 
         """
-        return theta.sigma
+        if isinstance(theta, GBMparam):
+            return theta.sigma
+        return float(theta[1])
 
     def betamat(self, theta: np.ndarray | Sequence[float]) -> np.ndarray:
         """Coefficients in linear representation of the first moment.
@@ -109,7 +113,7 @@ def gammamat(self, theta: np.ndarray | Sequence[float]) -> np.ndarray:
         scale = float(self.exact_scale(state=np.array(0), theta=param))
         return np.array([loc**2 + scale**2, 0], dtype=float)
 
-    def dbetamat(self, theta: np.ndarray | Sequence[float]) -> np.ndarray:
+    def dbetamat(self, theta: GenericParam | np.ndarray | Sequence[float]) -> np.ndarray:
         """Calculate derivative of the first moment coefficients (numerical).
 
         Parameters
@@ -126,7 +130,7 @@ def dbetamat(self, theta: np.ndarray | Sequence[float]) -> np.ndarray:
         with np.errstate(divide="ignore"):
             return nd.Jacobian(self.betamat)(theta)
 
-    def dgammamat(self, theta: np.ndarray | Sequence[float]) -> np.ndarray:
+    def dgammamat(self, theta: GenericParam | np.ndarray | Sequence[float]) -> np.ndarray:
         """Calculate derivative of the second moment coefficients (numerical).
 
         Parameters
@@ -188,7 +192,7 @@ def dgammamat_exact(self, theta: np.ndarray | Sequence[float]) -> np.ndarray:
         )
 
     @staticmethod
-    def realized_depvar(*, data: np.ndarray, subset: slice | None = None) -> np.ndarray:  # noqa: ARG004
+    def realized_depvar(*, data: np.ndarray | Sequence[np.ndarray], subset: slice | None = None) -> np.ndarray:  # noqa: ARG004
         """Array of the left-hand side variables in realized moment conditions.
 
         Parameters
@@ -210,8 +214,8 @@ def realized_depvar(*, data: np.ndarray, subset: slice | None = None) -> np.ndar
     def realized_const(
         self,
         *,
-        param: Any = None,  # noqa: ANN401
-        aggh: Any = 1,  # noqa: ARG002, ANN401
+        param: GenericParam | np.ndarray | Sequence[float] | None = None,
+        aggh: float = 1,  # noqa: ARG002
         subset: slice | None = None,  # noqa: ARG002
     ) -> np.ndarray:
         """Intercept in the realized moment conditions.
@@ -231,11 +235,19 @@ def realized_const(
             Intercept
 
         """
-        theta = param
-        mean, sigma = float(theta[0]), float(theta[1])
+        if param is None:
+            param = self.param
+        assert param is not None
+        if isinstance(param, GBMparam):
+            mean, sigma = param.mean, param.sigma
+        elif isinstance(param, (np.ndarray, Sequence)):
+            param_arr = np.asarray(param)
+            mean, sigma = float(param_arr[0]), float(param_arr[1])
+        else:
+            raise TypeError("Invalid param type for realized_const")
         return np.array([mean - sigma**2 / 2, sigma**2, sigma**4])
 
-    def drealized_const(self, theta: Any) -> np.ndarray:  # noqa: ANN401
+    def drealized_const(self, theta: GenericParam | np.ndarray | Sequence[float]) -> np.ndarray:
         """Calculate derivative of the intercept in the realized moment conditions.
 
         Parameters
@@ -250,14 +262,14 @@ def drealized_const(self, theta: Any) -> np.ndarray:  # noqa: ANN401
 
         """
 
-        def _realized_const_wrapper(theta: Any) -> np.ndarray:  # noqa: ANN401
+        def _realized_const_wrapper(theta: GenericParam | np.ndarray | Sequence[float]) -> np.ndarray:
             return self.realized_const(param=theta)
 
         with np.errstate(divide="ignore"):
             return nd.Jacobian(_realized_const_wrapper)(theta)
 
     @staticmethod
-    def instruments(*, data: Any, instrlag: int = 1) -> np.ndarray:  # noqa: ANN401
+    def instruments(*, data: np.ndarray | Sequence[np.ndarray], instrlag: int = 1) -> np.ndarray:
         """Create an array of instruments.
 
         Parameters
@@ -280,11 +292,11 @@ def instruments(*, data: Any, instrlag: int = 1) -> np.ndarray:  # noqa: ANN401
     def integrated_mom(
         self,
         *,
-        theta: Any,  # noqa: ANN401
-        data: Any = None,  # noqa: ANN401
-        instr_data: Any = None,  # noqa: ARG002, ANN401
+        theta: GenericParam | np.ndarray | Sequence[float],
+        data: np.ndarray | Sequence[np.ndarray] | None = None,
+        instr_data: np.ndarray | None = None,  # noqa: ARG002
         instr_choice: str = "const",  # noqa: ARG002
-        aggh: Any = 1,  # noqa: ARG002, ANN401
+        aggh: float | Sequence[float] = 1,  # noqa: ARG002
         subset: str = "all",  # noqa: ARG002
         instrlag: int = 1,
         measure: str = "P",  # noqa: ARG002
@@ -337,8 +349,8 @@ def integrated_mom(
     def momcond(
         self,
         *,
-        theta: np.ndarray | Sequence[float],
-        data: Any = None,  # noqa: ANN401
+        theta: GenericParam | np.ndarray | Sequence[float],
+        data: np.ndarray | Sequence[np.ndarray] | None = None,
         instrlag: int = 1,
     ) -> tuple[np.ndarray, np.ndarray]:
         """Moment function.
@@ -368,7 +380,11 @@ def momcond(
         datamat = np.hstack([np.ones((nobs, 1)), lagdata])
 
         # Coefficients in the first moment (mean)
-        linearcoef = [self.betamat(theta), self.gammamat(theta)]
+        if isinstance(theta, GenericParam):
+            theta_vec = theta.get_theta()
+        else:
+            theta_vec = theta
+        linearcoef = [self.betamat(theta_vec), self.gammamat(theta_vec)]
         # Coefficients in the second moment (variance)
         dlinearcoef = [self.dbetamat(theta), self.dgammamat(theta)]
 
diff --git a/src/affidiff/model_generic.py b/src/affidiff/model_generic.py
index 04c2c91..f961812 100644
--- a/src/affidiff/model_generic.py
+++ b/src/affidiff/model_generic.py
@@ -7,7 +7,7 @@
 from typing import TYPE_CHECKING, Any, Sequence, cast
 
 import numpy as np
-from mygmm import GMM
+from mygmm import GMM, Results
 
 from affidiff.helper_functions import ajd_diff, ajd_drift, columnwise_prod, instruments, nice_errors, rolling_window
 
@@ -43,7 +43,7 @@ class SDE(ABC):
 
     """
 
-    def __init__(self, param: Any = None) -> None:  # noqa: ANN401
+    def __init__(self, param: GenericParam | None = None) -> None:
         """Initialize the class.
 
         Parameters
@@ -74,21 +74,21 @@ def get_start(self) -> np.ndarray | list[float]:
         raise NotImplementedError("Must be overridden")
 
     @staticmethod
-    def realized_depvar(*, data: Any, subset: Any = None) -> Any:  # noqa: ANN401
+    def realized_depvar(*, data: np.ndarray | Sequence[np.ndarray], subset: slice | None = None) -> np.ndarray:
         """Realized dependent variables."""
         raise NotImplementedError("Must be overridden")
 
-    def mat_a(self, *, param: Any, subset: Any = None) -> Any:  # noqa: ANN401
+    def mat_a(self, *, param: GenericParam, subset: slice | None = None) -> np.ndarray:
         """Matrix A in integrated moments."""
         raise NotImplementedError("Must be overridden")
 
     def realized_const(
         self,
         *,
-        param: Any = None,  # noqa: ANN401
-        aggh: Any = 1,  # noqa: ANN401
-        subset: Any = None,  # noqa: ANN401
-    ) -> Any:  # noqa: ANN401
+        param: GenericParam | None = None,
+        aggh: float = 1,
+        subset: slice | None = None,
+    ) -> np.ndarray:
         """Realized constant in integrated moments."""
         raise NotImplementedError("Must be overridden")
 
@@ -173,30 +173,30 @@ def exact_scale(self, *, state: np.ndarray, theta: GenericParam) -> np.ndarray:
         return self.euler_scale(state=state, theta=theta)
 
     @staticmethod
-    def mean_vol(*, param: Any, aggh: float) -> float:  # noqa: ANN401
+    def mean_vol(*, param: GenericParam, aggh: float) -> float:
         """Unconditional mean of volatility."""
         raise NotImplementedError
 
     @staticmethod
-    def mean_vol2(*, param: Any, aggh: float) -> float:  # noqa: ANN401
+    def mean_vol2(*, param: GenericParam, aggh: float) -> float:
         """Unconditional mean of squared volatility."""
         raise NotImplementedError
 
     @staticmethod
-    def mean_ret(*, param: Any, aggh: float) -> float:  # noqa: ANN401
+    def mean_ret(*, param: GenericParam, aggh: float) -> float:
         """Unconditional mean of returns."""
         raise NotImplementedError
 
     @staticmethod
-    def mean_cross(*, param: Any, aggh: float) -> float:  # noqa: ANN401
+    def mean_cross(*, param: GenericParam, aggh: float) -> float:
         """Unconditional mean of returns times volatility."""
         raise NotImplementedError
 
     def momcond(
         self,
         *,
-        theta: Any,  # noqa: ANN401
-        data: Any = None,  # noqa: ANN401
+        theta: GenericParam | np.ndarray | Sequence[float],
+        data: np.ndarray | Sequence[np.ndarray] | None = None,
         instrlag: int = 1,
     ) -> tuple[np.ndarray, np.ndarray]:
         """Moment conditions."""
@@ -255,15 +255,15 @@ def update(self, *, state: np.ndarray, error: np.ndarray) -> np.ndarray:
     def simulate(
         self,
         *,
-        start: Any = None,  # noqa: ANN401
+        start: np.ndarray | float | Sequence[float] | None = None,
         nsub: int = 80,
         ndiscr: int = 1,
         nobs: int = 500,
         nsim: int = 1,
-        diff: Any = None,  # noqa: ANN401
+        diff: int | Sequence[int] | slice | None = None,
         new_innov: bool = True,
         cython: bool = False,
-    ) -> Any:  # noqa: ANN401
+    ) -> np.ndarray:
         """Simulate observations from the model.
 
         Parameters
@@ -338,16 +338,16 @@ def simulate(
     def sim_realized(
         self,
         *,
-        start: Any = None,  # noqa: ANN401
+        start: np.ndarray | float | Sequence[float] | None = None,
         nsub: int = 80,
         ndiscr: int = 10,
         aggh: int = 1,
         nperiods: int = 500,
         nsim: int = 1,
-        diff: Any = None,  # noqa: ANN401
+        diff: int | Sequence[int] | slice | None = None,
         new_innov: bool = True,
         cython: bool = False,
-    ) -> tuple[Any, Any]:  # noqa: ANN401
+    ) -> tuple[np.ndarray, np.ndarray]:
         """Simulate realized returns and variance from the model.
 
         Parameters
@@ -406,7 +406,7 @@ def sim_realized_pq(
         ndiscr: int = 10,
         nperiods: int = 500,
         nsim: int = 1,
-        diff: Any = None,  # noqa: ANN401
+        diff: int | Sequence[int] | slice | None = None,
         new_innov: bool = True,
         cython: bool = False,
     ) -> tuple[tuple[np.ndarray, np.ndarray], tuple[np.ndarray, np.ndarray]]:
@@ -477,14 +477,14 @@ def sim_realized_pq(
     def gmmest(
         self,
         *,
-        theta_start: Any,  # noqa: ANN401
-        data: Any = None,  # noqa: ANN401
+        theta_start: GenericParam,
+        data: np.ndarray | Sequence[np.ndarray] | None = None,
         instrlag: int = 1,
         iter: int = 2,
         method: str = "BFGS",
         kernel: str = "Bartlett",
         band: int | None = None,
-    ) -> Any:  # noqa: ANN401
+    ) -> Results:
         """Estimate model parameters using GMM.
 
         Parameters
@@ -524,8 +524,8 @@ def integrated_gmm(
         self,
         *,
         param_start: GenericParam,
-        data: Any = None,  # noqa: ANN401
-        instr_data: Any = None,  # noqa: ANN401
+        data: np.ndarray | Sequence[Any] | tuple[Any, ...] | None = None,
+        instr_data: np.ndarray | None = None,
         instr_choice: str = "const",
         aggh: float | Sequence[float] = 1,
         instrlag: int = 1,
@@ -533,12 +533,12 @@ def integrated_gmm(
         measure: str = "P",
         names: list[str] | None = None,
         bounds: list[tuple[float | None, float | None]] | None = None,
-        constraints: Any = (),  # noqa: ANN401
+        constraints: Sequence[dict[str, object]] | dict[str, object] | tuple[()] = (),
         iter: int = 2,
         method: str = "BFGS",
         kernel: str = "Bartlett",
         band: int | None = None,
-    ) -> Any:  # noqa: ANN401
+    ) -> Results:
         """Estimate model parameters using Integrated GMM.
 
         Parameters
@@ -585,12 +585,12 @@ def integrated_gmm(
 
         """
         estimator = GMM(self.integrated_mom)
-        self.param: Any = param_start
-        theta_start = self.param.get_theta(subset=subset, measure=measure)
+        self.param = param_start
+        theta_start = self.param.get_theta(subset=subset, measure=measure)  # type: ignore[call-arg]
         if names is None:
-            names = self.param.get_names(subset=subset, measure=measure)
+            names = self.param.get_names(subset=subset, measure=measure)  # type: ignore[call-arg]
         if bounds is None:
-            bounds = self.param.get_bounds(subset=subset, measure=measure)
+            bounds = self.param.get_bounds(subset=subset, measure=measure)  # type: ignore[call-arg]
         if constraints == ():
             constraints = self.param.get_constraints()
         return estimator.gmmest(
@@ -615,14 +615,14 @@ def integrated_mom(
         self,
         *,
         theta: np.ndarray | Sequence[float],
-        data: Any = None,  # noqa: ANN401
-        instr_data: Any = None,  # noqa: ANN401
+        data: np.ndarray | Sequence[np.ndarray] | None = None,
+        instr_data: np.ndarray | None = None,
         instr_choice: str = "const",
         aggh: float | Sequence[float] = 1,
         subset: str = "all",
         instrlag: int = 1,
         measure: str = "P",
-    ) -> tuple[np.ndarray, Any]:
+    ) -> tuple[np.ndarray, np.ndarray | None]:
         """Integrated moment function.
 
         Parameters
@@ -665,6 +665,7 @@ def integrated_mom(
         if subset == "vol":
             subset_sl = slice(2)
 
+        assert self.param is not None
         self.param.update(theta=theta, subset=subset, measure=measure)
         lag = 2
 
@@ -680,16 +681,18 @@ def integrated_mom(
                 # (nobs - lag, 4) array
                 error.append(
                     depvar.dot(self.mat_a(param=self.param, subset=subset_sl).T)
-                    - self.realized_const(param=self.param, aggh=agg, subset=subset_sl)
+                    - self.realized_const(param=self.param, aggh=float(cast(float, agg)), subset=subset_sl)
                 )
 
             error = np.hstack(error)
 
         else:
-            depvar = self.realized_depvar(data=data)[lag:]  # type: ignore[arg-type]
+            assert data is not None
+            depvar = self.realized_depvar(data=data)[lag:]
+            aggh_val = aggh[0] if isinstance(aggh, Sequence) and not isinstance(aggh, (str, bytes)) else float(aggh)  # type: ignore[arg-type]
             # (nobs - lag, 4) array
             error = depvar.dot(self.mat_a(param=self.param, subset=subset_sl).T) - self.realized_const(
-                param=self.param, aggh=aggh, subset=subset_sl
+                param=self.param, aggh=float(cast(float, aggh_val)), subset=subset_sl
             )
 
         nobs = error.shape[0] + lag
diff --git a/src/affidiff/model_heston.py b/src/affidiff/model_heston.py
index b96f1c3..e3192bf 100644
--- a/src/affidiff/model_heston.py
+++ b/src/affidiff/model_heston.py
@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, cast
+from typing import TYPE_CHECKING, Any, Sequence, cast
 
 import numpy as np
 from statsmodels.tsa.tsatools import lagmat
@@ -11,15 +11,15 @@
 from affidiff.param_heston import HestonParam
 
 if TYPE_CHECKING:
-    pass
+    from affidiff.param_generic import GenericParam
 
 
 class Heston(SDE):
     """Heston model."""
 
-    param: Any
+    param: HestonParam | None
 
-    def __init__(self, param: Any = None) -> None:  # noqa: ANN401
+    def __init__(self, param: HestonParam | None = None) -> None:
         """Initialize the class.
 
         Parameters
@@ -41,10 +41,11 @@ def get_start(self) -> list[float]:
             Starting values for price and variance
 
         """
+        assert self.param is not None
         return [1.0, float(self.param.mean_v)]
 
     @staticmethod
-    def coef_big_a(*, param: Any, aggh: float) -> float:  # noqa: ANN401
+    def coef_big_a(*, param: HestonParam, aggh: float) -> float:
         """Coefficient A_h in exact discretization of volatility.
 
         Parameters
@@ -61,7 +62,7 @@ def coef_big_a(*, param: Any, aggh: float) -> float:  # noqa: ANN401
         """
         return float(np.exp(-param.kappa * aggh))
 
-    def coef_big_c(self, *, param: Any, aggh: float) -> float:  # noqa: ANN401
+    def coef_big_c(self, *, param: HestonParam, aggh: float) -> float:
         """Coefficient C_h in exact discretization of volatility.
 
         Parameters
@@ -78,7 +79,7 @@ def coef_big_c(self, *, param: Any, aggh: float) -> float:  # noqa: ANN401
         """
         return float(param.mean_v * (1 - self.coef_big_a(param=param, aggh=aggh)))
 
-    def coef_small_a(self, *, param: Any, aggh: float) -> float:  # noqa: ANN401
+    def coef_small_a(self, *, param: HestonParam, aggh: float) -> float:
         """Coefficient a_h in exact discretization of volatility.
 
         Parameters
@@ -95,7 +96,7 @@ def coef_small_a(self, *, param: Any, aggh: float) -> float:  # noqa: ANN401
         """
         return float((1 - self.coef_big_a(param=param, aggh=aggh)) / param.kappa / aggh)
 
-    def coef_small_c(self, *, param: Any, aggh: float) -> float:  # noqa: ANN401
+    def coef_small_c(self, *, param: HestonParam, aggh: float) -> float:
         """Coefficient c_h in exact discretization of volatility.
 
         Parameters
@@ -113,7 +114,7 @@ def coef_small_c(self, *, param: Any, aggh: float) -> float:  # noqa: ANN401
         return float(param.mean_v * (1 - self.coef_small_a(param=param, aggh=aggh)))
 
     @staticmethod
-    def mean_vol(*, param: Any, aggh: float) -> float:  # noqa: ARG004, ANN401
+    def mean_vol(*, param: GenericParam, aggh: float) -> float:  # noqa: ARG004
         """Unconditional mean of realized volatiliy.
 
         Parameters
@@ -128,9 +129,10 @@ def mean_vol(*, param: Any, aggh: float) -> float:  # noqa: ARG004, ANN401
         float
 
         """
+        assert isinstance(param, HestonParam)
         return float(param.mean_v)
 
-    def mean_vol2(self, *, param: Any, aggh: float) -> float:  # noqa: ANN401
+    def mean_vol2(self, *, param: GenericParam, aggh: float) -> float:
         """Unconditional mean of squared realized volatiliy.
 
         Parameters
@@ -145,12 +147,13 @@ def mean_vol2(self, *, param: Any, aggh: float) -> float:  # noqa: ANN401
         float
 
         """
+        assert isinstance(param, HestonParam)
         return float(
             (param.eta / param.kappa) ** 2 * self.coef_small_c(param=param, aggh=aggh) / aggh + param.mean_v**2
         )
 
     @staticmethod
-    def mean_ret(*, param: Any, aggh: float) -> float:  # noqa: ARG004, ANN401
+    def mean_ret(*, param: GenericParam, aggh: float) -> float:  # noqa: ARG004
         """Unconditional mean of realized returns.
 
         Parameters
@@ -165,9 +168,10 @@ def mean_ret(*, param: Any, aggh: float) -> float:  # noqa: ARG004, ANN401
         float
 
         """
+        assert isinstance(param, HestonParam)
         return float((param.lmbd - 0.5) * param.mean_v)
 
-    def mean_cross(self, *, param: Any, aggh: float) -> float:  # noqa: ANN401
+    def mean_cross(self, *, param: GenericParam, aggh: float) -> float:
         """Unconditional mean of realized returns times volatility.
 
         Parameters
@@ -182,13 +186,16 @@ def mean_cross(self, *, param: Any, aggh: float) -> float:  # noqa: ANN401
         float
 
         """
+        assert isinstance(param, HestonParam)
         p = param
         val = (p.lmbd - 0.5) * self.mean_vol2(param=param, aggh=aggh) + (
             p.rho * p.eta / p.kappa * self.coef_small_c(param=param, aggh=aggh) / aggh
         )
         return float(val)
 
-    def realized_const(self, *, param: Any = None, aggh: float = 1, subset: slice | None = None) -> np.ndarray:  # noqa: ANN401
+    def realized_const(
+        self, *, param: GenericParam | None = None, aggh: float = 1, subset: slice | None = None
+    ) -> np.ndarray:
         """Intercept in the realized moment conditions.
 
         Parameters
@@ -206,6 +213,9 @@ def realized_const(self, *, param: Any = None, aggh: float = 1, subset: slice |
             Intercept
 
         """
+        if param is None:
+            param = self.param
+        assert param is not None
         res = (
             (self.mat_a0(param=param, aggh=1) + self.mat_a1(param=param, aggh=1) + self.mat_a2(param=param, aggh=1))
             * self.depvar_unc_mean(param=param, aggh=aggh)
@@ -282,7 +292,7 @@ def mat_a2(self, *, param: Any, aggh: float) -> np.ndarray:  # noqa: ARG002, ANN
         mat_a[3, 1] = (param.lmbd - 0.5) * self.coef_big_a(param=param, aggh=1)
         return mat_a
 
-    def mat_a(self, *, param: Any, subset: slice | None = None) -> np.ndarray:  # noqa: ANN401
+    def mat_a(self, *, param: GenericParam, subset: slice | None = None) -> np.ndarray:
         """Matrix A in integrated moments.
 
         Parameters
@@ -298,6 +308,7 @@ def mat_a(self, *, param: Any, subset: slice | None = None) -> np.ndarray:  # no
             Matrix A
 
         """
+        assert isinstance(param, HestonParam)
         mat_a_tuple = (
             self.mat_a0(param=param, aggh=1),
             self.mat_a1(param=param, aggh=1),
@@ -309,7 +320,7 @@ def mat_a(self, *, param: Any, subset: slice | None = None) -> np.ndarray:  # no
         return np.squeeze(res)
 
     @staticmethod
-    def realized_depvar(*, data: Any, subset: slice | None = None) -> np.ndarray:  # noqa: ANN401
+    def realized_depvar(*, data: np.ndarray | Sequence[np.ndarray], subset: slice | None = None) -> np.ndarray:
         """Array of the left-hand side variables in realized moment conditions.
 
         Parameters
diff --git a/src/affidiff/model_vasicek.py b/src/affidiff/model_vasicek.py
index 212003d..9620a8f 100644
--- a/src/affidiff/model_vasicek.py
+++ b/src/affidiff/model_vasicek.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_vasicek import VasicekParam
+
 
 class Vasicek(SDE):
     """Vasicek model."""
 
-    def __init__(self, param: Any = None) -> None:  # noqa: ANN401
+    param: VasicekParam | None
+
+    def __init__(self, param: VasicekParam | 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: VasicekParam) -> 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) -> float:  # noqa: ARG004, ANN401
+    def diff(*, state: np.ndarray | float, theta: VasicekParam) -> float:  # noqa: ARG004
         """Diffusion (instantaneous volatility) function.
 
         Parameters
diff --git a/src/affidiff/param_generic.py b/src/affidiff/param_generic.py
index 1828a0c..77329e6 100644
--- a/src/affidiff/param_generic.py
+++ b/src/affidiff/param_generic.py
@@ -22,6 +22,10 @@ class GenericParam(ABC):
     """
 
     measure: str = "P"
+    mat_k0: Any
+    mat_k1: Any
+    mat_h0: Any
+    mat_h1: Any
 
     def __init__(self) -> None:
         """Initialize class."""
@@ -87,7 +91,7 @@ def get_model_name() -> str:
 
     @staticmethod
     @abstractmethod
-    def get_names() -> list[str]:
+    def get_names(*, subset: str = "all", measure: str = "PQ") -> list[str]:
         """Return parameter names.
 
         Returns
@@ -99,7 +103,7 @@ def get_names() -> list[str]:
         raise NotImplementedError("Must be overridden")
 
     @abstractmethod
-    def get_theta(self) -> np.ndarray:
+    def get_theta(self, *, subset: str = "all", measure: str = "P") -> np.ndarray:
         """Return vector of parameters.
 
         Returns

From f8272ab72a7ff78cc0849371a1bc983928ce084b Mon Sep 17 00:00:00 2001
From: Stanislav Khrapov <khrapovs@gmail.com>
Date: Thu, 30 Jul 2026 21:08:34 +0200
Subject: [PATCH 4/7] refactor: replace unused argument noqa tags with explicit
 underscore assignments across model and parameter modules

---
 src/affidiff/model_ct.py      | 6 ++++--
 src/affidiff/model_gbm.py     | 9 ++++++---
 src/affidiff/model_heston.py  | 9 ++++++---
 src/affidiff/model_vasicek.py | 3 ++-
 src/affidiff/param_cir.py     | 3 ++-
 src/affidiff/param_gbm.py     | 3 ++-
 src/affidiff/param_generic.py | 3 ++-
 src/affidiff/param_vasicek.py | 3 ++-
 8 files changed, 26 insertions(+), 13 deletions(-)

diff --git a/src/affidiff/model_ct.py b/src/affidiff/model_ct.py
index 842c262..164f599 100644
--- a/src/affidiff/model_ct.py
+++ b/src/affidiff/model_ct.py
@@ -247,7 +247,7 @@ def roots(self, *, param: CentTendParam, aggh: float) -> list[float]:
         ]
 
     @staticmethod
-    def mean_vol(*, param: GenericParam, aggh: float) -> float:  # noqa: ARG004
+    def mean_vol(*, param: GenericParam, aggh: float) -> float:
         """Unconditional mean of realized volatiliy.
 
         Parameters
@@ -262,6 +262,7 @@ def mean_vol(*, param: GenericParam, aggh: float) -> float:  # noqa: ARG004
         float
 
         """
+        _ = aggh
         assert isinstance(param, CentTendParam)
         return float(param.mean_v)
 
@@ -288,7 +289,7 @@ def mean_vol2(self, *, param: GenericParam, aggh: float) -> float:
         )
 
     @staticmethod
-    def mean_ret(*, param: GenericParam, aggh: float) -> float:  # noqa: ARG004
+    def mean_ret(*, param: GenericParam, aggh: float) -> float:
         """Unconditional mean of realized returns.
 
         Parameters
@@ -303,6 +304,7 @@ def mean_ret(*, param: GenericParam, aggh: float) -> float:  # noqa: ARG004
         float
 
         """
+        _ = aggh
         assert isinstance(param, CentTendParam)
         return float((param.lmbd - 0.5) * param.mean_v)
 
diff --git a/src/affidiff/model_gbm.py b/src/affidiff/model_gbm.py
index 6a862c9..f756b31 100644
--- a/src/affidiff/model_gbm.py
+++ b/src/affidiff/model_gbm.py
@@ -35,7 +35,7 @@ def get_start(self) -> list[float]:
         return [1.0]
 
     @staticmethod
-    def drift(*, state: np.ndarray | float, theta: GBMparam | np.ndarray | Sequence[float]) -> np.ndarray | float:  # noqa: ARG004
+    def drift(*, state: np.ndarray | float, theta: GBMparam | np.ndarray | Sequence[float]) -> np.ndarray | float:
         """Drift function.
 
         Parameters
@@ -51,12 +51,13 @@ def drift(*, state: np.ndarray | float, theta: GBMparam | np.ndarray | Sequence[
             Drift value
 
         """
+        _ = state
         if isinstance(theta, GBMparam):
             return theta.mean - theta.sigma**2 / 2
         return float(theta[0]) - float(theta[1]) ** 2 / 2
 
     @staticmethod
-    def diff(*, state: np.ndarray | float, theta: GBMparam | np.ndarray | Sequence[float]) -> np.ndarray | float:  # noqa: ARG004
+    def diff(*, state: np.ndarray | float, theta: GBMparam | np.ndarray | Sequence[float]) -> np.ndarray | float:
         """Diffusion (instantaneous volatility) function.
 
         Parameters
@@ -72,6 +73,7 @@ def diff(*, state: np.ndarray | float, theta: GBMparam | np.ndarray | Sequence[f
             Diffusion value
 
         """
+        _ = state
         if isinstance(theta, GBMparam):
             return theta.sigma
         return float(theta[1])
@@ -192,7 +194,7 @@ def dgammamat_exact(self, theta: np.ndarray | Sequence[float]) -> np.ndarray:
         )
 
     @staticmethod
-    def realized_depvar(*, data: np.ndarray | Sequence[np.ndarray], subset: slice | None = None) -> np.ndarray:  # noqa: ARG004
+    def realized_depvar(*, data: np.ndarray | Sequence[np.ndarray], subset: slice | None = None) -> np.ndarray:
         """Array of the left-hand side variables in realized moment conditions.
 
         Parameters
@@ -208,6 +210,7 @@ def realized_depvar(*, data: np.ndarray | Sequence[np.ndarray], subset: slice |
             Dependend variables
 
         """
+        _ = subset
         ret, rvar = data
         return np.vstack([ret, rvar, rvar**2])
 
diff --git a/src/affidiff/model_heston.py b/src/affidiff/model_heston.py
index e3192bf..cbd2331 100644
--- a/src/affidiff/model_heston.py
+++ b/src/affidiff/model_heston.py
@@ -114,7 +114,7 @@ def coef_small_c(self, *, param: HestonParam, aggh: float) -> float:
         return float(param.mean_v * (1 - self.coef_small_a(param=param, aggh=aggh)))
 
     @staticmethod
-    def mean_vol(*, param: GenericParam, aggh: float) -> float:  # noqa: ARG004
+    def mean_vol(*, param: GenericParam, aggh: float) -> float:
         """Unconditional mean of realized volatiliy.
 
         Parameters
@@ -129,6 +129,7 @@ def mean_vol(*, param: GenericParam, aggh: float) -> float:  # noqa: ARG004
         float
 
         """
+        _ = aggh
         assert isinstance(param, HestonParam)
         return float(param.mean_v)
 
@@ -153,7 +154,7 @@ def mean_vol2(self, *, param: GenericParam, aggh: float) -> float:
         )
 
     @staticmethod
-    def mean_ret(*, param: GenericParam, aggh: float) -> float:  # noqa: ARG004
+    def mean_ret(*, param: GenericParam, aggh: float) -> float:
         """Unconditional mean of realized returns.
 
         Parameters
@@ -168,6 +169,7 @@ def mean_ret(*, param: GenericParam, aggh: float) -> float:  # noqa: ARG004
         float
 
         """
+        _ = aggh
         assert isinstance(param, HestonParam)
         return float((param.lmbd - 0.5) * param.mean_v)
 
@@ -225,7 +227,7 @@ def realized_const(
         return np.squeeze(res)
 
     @staticmethod
-    def mat_a0(*, param: Any, aggh: float) -> np.ndarray:  # noqa: ARG004, ANN401
+    def mat_a0(*, param: Any, aggh: float) -> np.ndarray:  # noqa: ANN401
         """Matrix A_0 in integrated moments.
 
         Parameters
@@ -241,6 +243,7 @@ def mat_a0(*, param: Any, aggh: float) -> np.ndarray:  # noqa: ARG004, ANN401
             Matrix A_0
 
         """
+        _ = (param, aggh)
         return np.diag([0, 1, 0, 0]).astype(float)
 
     def mat_a1(self, *, param: Any, aggh: float) -> np.ndarray:  # noqa: ARG002, ANN401
diff --git a/src/affidiff/model_vasicek.py b/src/affidiff/model_vasicek.py
index 9620a8f..49d076d 100644
--- a/src/affidiff/model_vasicek.py
+++ b/src/affidiff/model_vasicek.py
@@ -60,7 +60,7 @@ def drift(*, state: np.ndarray | float, theta: VasicekParam) -> np.ndarray | flo
         return theta.kappa * (theta.mean - state)
 
     @staticmethod
-    def diff(*, state: np.ndarray | float, theta: VasicekParam) -> float:  # noqa: ARG004
+    def diff(*, state: np.ndarray | float, theta: VasicekParam) -> float:
         """Diffusion (instantaneous volatility) function.
 
         Parameters
@@ -76,6 +76,7 @@ def diff(*, state: np.ndarray | float, theta: VasicekParam) -> float:  # noqa: A
             Diffusion value
 
         """
+        _ = state
         return theta.eta
 
 
diff --git a/src/affidiff/param_cir.py b/src/affidiff/param_cir.py
index 0577cae..f6b1008 100644
--- a/src/affidiff/param_cir.py
+++ b/src/affidiff/param_cir.py
@@ -117,7 +117,7 @@ def get_model_name() -> str:
         return "CIR"
 
     @staticmethod
-    def get_names(*, subset: str = "all", measure: str = "PQ") -> list[str]:  # noqa: ARG004
+    def get_names(*, subset: str = "all", measure: str = "PQ") -> list[str]:
         """Return parameter names.
 
         Returns
@@ -126,6 +126,7 @@ def get_names(*, subset: str = "all", measure: str = "PQ") -> list[str]:  # noqa
             Parameter names
 
         """
+        _ = (subset, measure)
         return ["mean", "kappa", "eta"]
 
     def get_theta(self, *, subset: str = "all", measure: str = "PQ") -> np.ndarray:  # noqa: ARG002
diff --git a/src/affidiff/param_gbm.py b/src/affidiff/param_gbm.py
index f82d663..f3bb701 100644
--- a/src/affidiff/param_gbm.py
+++ b/src/affidiff/param_gbm.py
@@ -110,7 +110,7 @@ def get_model_name() -> str:
         return "GBM"
 
     @staticmethod
-    def get_names(*, subset: str = "all", measure: str = "PQ") -> list[str]:  # noqa: ARG004
+    def get_names(*, subset: str = "all", measure: str = "PQ") -> list[str]:
         """Return parameter names.
 
         Returns
@@ -119,6 +119,7 @@ def get_names(*, subset: str = "all", measure: str = "PQ") -> list[str]:  # noqa
             Parameter names
 
         """
+        _ = (subset, measure)
         return ["mean", "sigma"]
 
     def get_theta(self, *, subset: str = "all", measure: str = "PQ") -> np.ndarray:  # noqa: ARG002
diff --git a/src/affidiff/param_generic.py b/src/affidiff/param_generic.py
index 77329e6..7d961bd 100644
--- a/src/affidiff/param_generic.py
+++ b/src/affidiff/param_generic.py
@@ -115,7 +115,7 @@ def get_theta(self, *, subset: str = "all", measure: str = "P") -> np.ndarray:
         raise NotImplementedError("Must be overridden")
 
     @staticmethod
-    def get_bounds(*, subset: str = "all", measure: str = "PQ") -> list[tuple[float | None, float | None]] | None:  # noqa: ARG004
+    def get_bounds(*, subset: str = "all", measure: str = "PQ") -> list[tuple[float | None, float | None]] | None:
         """Get parameter bounds.
 
         Returns
@@ -124,6 +124,7 @@ def get_bounds(*, subset: str = "all", measure: str = "PQ") -> list[tuple[float
             Parameter bounds
 
         """
+        _ = (subset, measure)
         return None
 
     def get_constraints(self) -> tuple[dict[str, Any], ...] | list[dict[str, Any]] | tuple[()]:
diff --git a/src/affidiff/param_vasicek.py b/src/affidiff/param_vasicek.py
index 2ac7e2f..bfd7243 100644
--- a/src/affidiff/param_vasicek.py
+++ b/src/affidiff/param_vasicek.py
@@ -115,7 +115,7 @@ def get_model_name() -> str:
         return "Vasicek"
 
     @staticmethod
-    def get_names(*, subset: str = "all", measure: str = "PQ") -> list[str]:  # noqa: ARG004
+    def get_names(*, subset: str = "all", measure: str = "PQ") -> list[str]:
         """Return parameter names.
 
         Returns
@@ -124,6 +124,7 @@ def get_names(*, subset: str = "all", measure: str = "PQ") -> list[str]:  # noqa
             Parameter names
 
         """
+        _ = (subset, measure)
         return ["mean", "kappa", "eta"]
 
     def get_theta(self, *, subset: str = "all", measure: str = "PQ") -> np.ndarray:  # noqa: ARG002

From 8311a2054faee10c5b62182cd58acee784683b12 Mon Sep 17 00:00:00 2001
From: Stanislav Khrapov <khrapovs@gmail.com>
Date: Thu, 30 Jul 2026 21:09:51 +0200
Subject: [PATCH 5/7] refactor: replace inline noqa: ARG002 comments with
 explicit underscore assignments for unused parameters.

---
 src/affidiff/model_gbm.py     | 16 +++++++++-------
 src/affidiff/model_heston.py  |  6 ++++--
 src/affidiff/param_cir.py     |  9 ++++++---
 src/affidiff/param_gbm.py     |  9 ++++++---
 src/affidiff/param_vasicek.py |  9 ++++++---
 5 files changed, 31 insertions(+), 18 deletions(-)

diff --git a/src/affidiff/model_gbm.py b/src/affidiff/model_gbm.py
index f756b31..6077486 100644
--- a/src/affidiff/model_gbm.py
+++ b/src/affidiff/model_gbm.py
@@ -218,8 +218,8 @@ def realized_const(
         self,
         *,
         param: GenericParam | np.ndarray | Sequence[float] | None = None,
-        aggh: float = 1,  # noqa: ARG002
-        subset: slice | None = None,  # noqa: ARG002
+        aggh: float = 1,
+        subset: slice | None = None,
     ) -> np.ndarray:
         """Intercept in the realized moment conditions.
 
@@ -238,6 +238,7 @@ def realized_const(
             Intercept
 
         """
+        _ = (aggh, subset)
         if param is None:
             param = self.param
         assert param is not None
@@ -297,12 +298,12 @@ def integrated_mom(
         *,
         theta: GenericParam | np.ndarray | Sequence[float],
         data: np.ndarray | Sequence[np.ndarray] | None = None,
-        instr_data: np.ndarray | None = None,  # noqa: ARG002
-        instr_choice: str = "const",  # noqa: ARG002
-        aggh: float | Sequence[float] = 1,  # noqa: ARG002
-        subset: str = "all",  # noqa: ARG002
+        instr_data: np.ndarray | None = None,
+        instr_choice: str = "const",
+        aggh: float | Sequence[float] = 1,
+        subset: str = "all",
         instrlag: int = 1,
-        measure: str = "P",  # noqa: ARG002
+        measure: str = "P",
     ) -> tuple[np.ndarray, np.ndarray]:
         """Integrated moment function.
 
@@ -333,6 +334,7 @@ def integrated_mom(
             Average derivative of the moment restrictions
 
         """
+        _ = (instr_data, instr_choice, aggh, subset, measure)
         assert data is not None
         # (nobs - instrlag, 3) array
         error = self.realized_depvar(data=data).T[instrlag:] - self.realized_const(param=theta)
diff --git a/src/affidiff/model_heston.py b/src/affidiff/model_heston.py
index cbd2331..30c50ca 100644
--- a/src/affidiff/model_heston.py
+++ b/src/affidiff/model_heston.py
@@ -246,7 +246,7 @@ def mat_a0(*, param: Any, aggh: float) -> np.ndarray:  # noqa: ANN401
         _ = (param, aggh)
         return np.diag([0, 1, 0, 0]).astype(float)
 
-    def mat_a1(self, *, param: Any, aggh: float) -> np.ndarray:  # noqa: ARG002, ANN401
+    def mat_a1(self, *, param: Any, aggh: float) -> np.ndarray:  # noqa: ANN401
         """Matrix A_1 in integrated moments.
 
         Parameters
@@ -262,12 +262,13 @@ def mat_a1(self, *, param: Any, aggh: float) -> np.ndarray:  # noqa: ARG002, ANN
             Matrix A_1
 
         """
+        _ = aggh
         mat_a = np.diag([1, 0, 0, 1]).astype(float)
         mat_a[1, 1] = -self.coef_big_a(param=param, aggh=1) * (1 + self.coef_big_a(param=param, aggh=1))
         mat_a[3, 1] = 0.5 - param.lmbd
         return mat_a
 
-    def mat_a2(self, *, param: Any, aggh: float) -> np.ndarray:  # noqa: ARG002, ANN401
+    def mat_a2(self, *, param: Any, aggh: float) -> np.ndarray:  # noqa: ANN401
         """Matrix A_2 in integrated moments.
 
         Parameters
@@ -283,6 +284,7 @@ def mat_a2(self, *, param: Any, aggh: float) -> np.ndarray:  # noqa: ARG002, ANN
             Matrix A_2
 
         """
+        _ = aggh
         mat_a = np.diag(
             [
                 -self.coef_big_a(param=param, aggh=1),
diff --git a/src/affidiff/param_cir.py b/src/affidiff/param_cir.py
index f6b1008..83af836 100644
--- a/src/affidiff/param_cir.py
+++ b/src/affidiff/param_cir.py
@@ -28,7 +28,7 @@ class CIRparam(GenericParam):
 
     """
 
-    def __init__(self, *, mean: float = 0.5, kappa: float = 1.5, eta: float = 0.1, measure: str = "P") -> None:  # noqa: ARG002
+    def __init__(self, *, mean: float = 0.5, kappa: float = 1.5, eta: float = 0.1, measure: str = "P") -> None:
         """Initialize class.
 
         Parameters
@@ -46,6 +46,7 @@ def __init__(self, *, mean: float = 0.5, kappa: float = 1.5, eta: float = 0.1, m
                 - 'Q' : risk-neutral
 
         """
+        _ = measure
         super().__init__()
         self.mean = mean
         self.kappa = kappa
@@ -88,7 +89,7 @@ def from_theta(cls, *, theta: np.ndarray | Sequence[float]) -> Self:
         param.update_ajd()
         return param
 
-    def update(self, *, theta: np.ndarray | Sequence[float], subset: str = "all", measure: str = "P") -> None:  # noqa: ARG002
+    def update(self, *, theta: np.ndarray | Sequence[float], subset: str = "all", measure: str = "P") -> None:
         """Update attributes from parameter vector.
 
         Parameters
@@ -101,6 +102,7 @@ def update(self, *, theta: np.ndarray | Sequence[float], subset: str = "all", me
             Probability measure
 
         """
+        _ = (subset, measure)
         self.mean, self.kappa, self.eta = float(theta[0]), float(theta[1]), float(theta[2])
         self.update_ajd()
 
@@ -129,7 +131,7 @@ def get_names(*, subset: str = "all", measure: str = "PQ") -> list[str]:
         _ = (subset, measure)
         return ["mean", "kappa", "eta"]
 
-    def get_theta(self, *, subset: str = "all", measure: str = "PQ") -> np.ndarray:  # noqa: ARG002
+    def get_theta(self, *, subset: str = "all", measure: str = "PQ") -> np.ndarray:
         """Return vector of parameters.
 
         Returns
@@ -138,4 +140,5 @@ def get_theta(self, *, subset: str = "all", measure: str = "PQ") -> np.ndarray:
             Parameter vector
 
         """
+        _ = (subset, measure)
         return np.array([self.mean, self.kappa, self.eta])
diff --git a/src/affidiff/param_gbm.py b/src/affidiff/param_gbm.py
index f3bb701..b9ce405 100644
--- a/src/affidiff/param_gbm.py
+++ b/src/affidiff/param_gbm.py
@@ -26,7 +26,7 @@ class GBMparam(GenericParam):
 
     """
 
-    def __init__(self, *, mean: float = 0.0, sigma: float = 0.2, measure: str = "P") -> None:  # noqa: ARG002
+    def __init__(self, *, mean: float = 0.0, sigma: float = 0.2, measure: str = "P") -> None:
         """Initialize class.
 
         Parameters
@@ -42,6 +42,7 @@ def __init__(self, *, mean: float = 0.0, sigma: float = 0.2, measure: str = "P")
                 - 'Q' : risk-neutral
 
         """
+        _ = measure
         super().__init__()
         self.mean = mean
         self.sigma = sigma
@@ -81,7 +82,7 @@ def from_theta(cls, *, theta: np.ndarray | Sequence[float]) -> Self:
         param.update_ajd()
         return param
 
-    def update(self, *, theta: np.ndarray | Sequence[float], subset: str = "all", measure: str = "P") -> None:  # noqa: ARG002
+    def update(self, *, theta: np.ndarray | Sequence[float], subset: str = "all", measure: str = "P") -> None:
         """Update attributes from parameter vector.
 
         Parameters
@@ -94,6 +95,7 @@ def update(self, *, theta: np.ndarray | Sequence[float], subset: str = "all", me
             Probability measure
 
         """
+        _ = (subset, measure)
         self.mean, self.sigma = float(theta[0]), float(theta[1])
         self.update_ajd()
 
@@ -122,7 +124,7 @@ def get_names(*, subset: str = "all", measure: str = "PQ") -> list[str]:
         _ = (subset, measure)
         return ["mean", "sigma"]
 
-    def get_theta(self, *, subset: str = "all", measure: str = "PQ") -> np.ndarray:  # noqa: ARG002
+    def get_theta(self, *, subset: str = "all", measure: str = "PQ") -> np.ndarray:
         """Return vector of parameters.
 
         Returns
@@ -131,4 +133,5 @@ def get_theta(self, *, subset: str = "all", measure: str = "PQ") -> np.ndarray:
             Parameter vector
 
         """
+        _ = (subset, measure)
         return np.array([self.mean, self.sigma])
diff --git a/src/affidiff/param_vasicek.py b/src/affidiff/param_vasicek.py
index bfd7243..2b63b97 100644
--- a/src/affidiff/param_vasicek.py
+++ b/src/affidiff/param_vasicek.py
@@ -28,7 +28,7 @@ class VasicekParam(GenericParam):
 
     """
 
-    def __init__(self, *, mean: float = 0.5, kappa: float = 1.5, eta: float = 0.1, measure: str = "P") -> None:  # noqa: ARG002
+    def __init__(self, *, mean: float = 0.5, kappa: float = 1.5, eta: float = 0.1, measure: str = "P") -> None:
         """Initialize class.
 
         Parameters
@@ -46,6 +46,7 @@ def __init__(self, *, mean: float = 0.5, kappa: float = 1.5, eta: float = 0.1, m
                 - 'Q' : risk-neutral
 
         """
+        _ = measure
         super().__init__()
         self.mean = mean
         self.kappa = kappa
@@ -86,7 +87,7 @@ def from_theta(cls, *, theta: np.ndarray | Sequence[float]) -> Self:
         param.update_ajd()
         return param
 
-    def update(self, *, theta: np.ndarray | Sequence[float], subset: str = "all", measure: str = "P") -> None:  # noqa: ARG002
+    def update(self, *, theta: np.ndarray | Sequence[float], subset: str = "all", measure: str = "P") -> None:
         """Update attributes from parameter vector.
 
         Parameters
@@ -99,6 +100,7 @@ def update(self, *, theta: np.ndarray | Sequence[float], subset: str = "all", me
             Probability measure
 
         """
+        _ = (subset, measure)
         self.mean, self.kappa, self.eta = float(theta[0]), float(theta[1]), float(theta[2])
         self.update_ajd()
 
@@ -127,7 +129,7 @@ def get_names(*, subset: str = "all", measure: str = "PQ") -> list[str]:
         _ = (subset, measure)
         return ["mean", "kappa", "eta"]
 
-    def get_theta(self, *, subset: str = "all", measure: str = "PQ") -> np.ndarray:  # noqa: ARG002
+    def get_theta(self, *, subset: str = "all", measure: str = "PQ") -> np.ndarray:
         """Return vector of parameters.
 
         Returns
@@ -136,4 +138,5 @@ def get_theta(self, *, subset: str = "all", measure: str = "PQ") -> np.ndarray:
             Parameter vector
 
         """
+        _ = (subset, measure)
         return np.array([self.mean, self.kappa, self.eta])

From 92b75667845fbfc3fbfaefc8fdfed1880baac52e Mon Sep 17 00:00:00 2001
From: Stanislav Khrapov <khrapovs@gmail.com>
Date: Thu, 30 Jul 2026 21:13:09 +0200
Subject: [PATCH 6/7] refactor: add type hints and import Results to Heston
 model simulation functions

---
 examples/try_centtend.py | 18 +++++++-----------
 examples/try_heston.py   | 25 ++++++++++++++-----------
 2 files changed, 21 insertions(+), 22 deletions(-)

diff --git a/examples/try_centtend.py b/examples/try_centtend.py
index a4ce2e4..ba21d76 100644
--- a/examples/try_centtend.py
+++ b/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.
 
     """
diff --git a/examples/try_heston.py b/examples/try_heston.py
index e809d49..900bcb5 100644
--- a/examples/try_heston.py
+++ b/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.
 
     """

From ae09f768d2808ef12e59c65c1cfa97558725a923 Mon Sep 17 00:00:00 2001
From: Stanislav Khrapov <khrapovs@gmail.com>
Date: Thu, 30 Jul 2026 21:33:12 +0200
Subject: [PATCH 7/7] refactor: enforce HestonParam type hints and validation
 in model methods

---
 src/affidiff/model_heston.py | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/src/affidiff/model_heston.py b/src/affidiff/model_heston.py
index 30c50ca..5312733 100644
--- a/src/affidiff/model_heston.py
+++ b/src/affidiff/model_heston.py
@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Sequence, cast
+from typing import TYPE_CHECKING, Sequence, cast
 
 import numpy as np
 from statsmodels.tsa.tsatools import lagmat
@@ -218,6 +218,7 @@ def realized_const(
         if param is None:
             param = self.param
         assert param is not None
+        assert isinstance(param, HestonParam)
         res = (
             (self.mat_a0(param=param, aggh=1) + self.mat_a1(param=param, aggh=1) + self.mat_a2(param=param, aggh=1))
             * self.depvar_unc_mean(param=param, aggh=aggh)
@@ -227,7 +228,7 @@ def realized_const(
         return np.squeeze(res)
 
     @staticmethod
-    def mat_a0(*, param: Any, aggh: float) -> np.ndarray:  # noqa: ANN401
+    def mat_a0(*, param: HestonParam, aggh: float) -> np.ndarray:
         """Matrix A_0 in integrated moments.
 
         Parameters
@@ -246,7 +247,7 @@ def mat_a0(*, param: Any, aggh: float) -> np.ndarray:  # noqa: ANN401
         _ = (param, aggh)
         return np.diag([0, 1, 0, 0]).astype(float)
 
-    def mat_a1(self, *, param: Any, aggh: float) -> np.ndarray:  # noqa: ANN401
+    def mat_a1(self, *, param: HestonParam, aggh: float) -> np.ndarray:
         """Matrix A_1 in integrated moments.
 
         Parameters
@@ -268,7 +269,7 @@ def mat_a1(self, *, param: Any, aggh: float) -> np.ndarray:  # noqa: ANN401
         mat_a[3, 1] = 0.5 - param.lmbd
         return mat_a
 
-    def mat_a2(self, *, param: Any, aggh: float) -> np.ndarray:  # noqa: ANN401
+    def mat_a2(self, *, param: HestonParam, aggh: float) -> np.ndarray:
         """Matrix A_2 in integrated moments.
 
         Parameters