test: fix linting and some tests [WIP]

Author: AldoGl

README.md

@@ -127,7 +127,7 @@ When the calibration terminates (~half a minute), towards the end  of the output
 you should see the following messages:
 ```
 True parameters:       [0.2, 0.2, 0.75]
-Best parameters found: [0.19 0.21 0.68]
+Best parameters found: [0.21 0.19 0.76]
 ```
 
 ## Docs

black_it/calibrator.py

@@ -128,9 +128,11 @@ def __init__(  # noqa: PLR0913
         # initialize arrays
         self.params_samp = np.zeros((0, self.param_grid.dims))
         self.losses_samp = np.zeros(0)
-        self.batch_num_samp = np.zeros(0, dtype=int)
-        self.method_samp = np.zeros(0, dtype=int)
-        self.series_samp = np.zeros((0, self.ensemble_size, self.N, self.D))
+        self.batch_num_samp: NDArray[np.int64] = np.zeros(0, dtype=int)
+        self.method_samp: NDArray[np.int64] = np.zeros(0, dtype=int)
+        self.series_samp: NDArray[np.float64] = np.zeros(
+            (0, self.ensemble_size, self.N, self.D),
+        )
 
         # initialize variables before calibration
         self.n_sampled_params = 0

black_it/loss_functions/gsl_div.py

@@ -270,7 +270,7 @@ def get_words(time_series: NDArray[np.float64], length: int) -> NDArray:
             "the chosen word length is too high",
             exception_class=ValueError,
         )
-        tsw = np.zeros(shape=(tswlen,), dtype=np.int32)
+        tsw: NDArray[np.float64] = np.zeros(shape=(tswlen,), dtype=np.int32)
 
         for i in range(length):
             k = 10 ** (length - i - 1)

black_it/loss_functions/likelihood.py

@@ -18,7 +18,7 @@
 from __future__ import annotations
 
 import warnings
-from typing import TYPE_CHECKING, Callable
+from typing import TYPE_CHECKING, Callable, cast
 
 import numpy as np
 
@@ -82,9 +82,13 @@ def compute_loss(
         Returns:
             The loss value.
         """
-        r = sim_data_ensemble.shape[0]  # number of repetitions
-        s = sim_data_ensemble.shape[1]  # simulation length
-        d = sim_data_ensemble.shape[2]  # number of dimensions
+        sim_data_ensemble_shape: tuple[int, int, int] = cast(
+            tuple[int, int, int],
+            sim_data_ensemble.shape,
+        )
+        r = sim_data_ensemble_shape[0]  # number of repetitions
+        s = sim_data_ensemble_shape[1]  # simulation length
+        d = sim_data_ensemble_shape[2]  # time series dimension
 
         if self.coordinate_weights is not None:
             warnings.warn(  # noqa: B028

black_it/plot/plot_results.py

@@ -33,6 +33,8 @@
     import os
     from collections.abc import Collection
 
+    from numpy.typing import NDArray
+
 
 def _get_samplers_id_table(saving_folder: str | os.PathLike) -> dict[str, int]:
     """Get the id table of the samplers from the checkpoint.
@@ -299,7 +301,7 @@ def plot_sampling_interact(saving_folder: str | os.PathLike) -> None:
     data_frame = pd.read_csv(calibration_results_file)
 
     max_bn = int(max(data_frame["batch_num_samp"]))
-    all_bns = np.arange(max_bn + 1, dtype=int)
+    all_bns: NDArray[np.int64] = np.arange(max_bn + 1, dtype=int)
     indices_bns = np.array_split(all_bns, min(max_bn, 3))
 
     dict_bns = {}

black_it/samplers/xgboost.py

@@ -28,9 +28,9 @@
 if TYPE_CHECKING:
     from numpy.typing import NDArray
 
-MAX_FLOAT32 = np.finfo(np.float32).max
-MIN_FLOAT32 = np.finfo(np.float32).min
-EPS_FLOAT32 = np.finfo(np.float32).eps
+MAX_FLOAT32: np.float64 = np.finfo(np.float32).max
+MIN_FLOAT32: np.float64 = np.finfo(np.float32).min
+EPS_FLOAT32: np.float64 = np.finfo(np.float32).eps
 
 
 class XGBoostSampler(MLSurrogateSampler):

black_it/search_space.py

@@ -72,7 +72,7 @@ def __init__(
         self._param_grid: list[NDArray[np.float64]] = []
         self._space_size = 1
         for i in range(self.dims):
-            new_col = np.arange(
+            new_col: NDArray[np.float64] = np.arange(
                 parameters_bounds[0][i],
                 parameters_bounds[1][i] + 0.0000001,
                 parameters_precision[i],

tests/fixtures/data/test_sir_w_breaks_python.npy

@@ -24,7 +24,7 @@
 
 
 TRUE_PARAMETERS_STR = "True parameters:       [0.2, 0.2, 0.75]"
-BEST_PARAMETERS_STR = "Best parameters found: [0.19 0.21 0.68]"
+BEST_PARAMETERS_STR = "Best parameters found: [0.21 0.19 0.76]"
 
 
 class TestMainExample(BaseMainExampleTestClass):

tests/test_examples/test_main.py

@@ -92,5 +92,5 @@ def test_sir_w_breaks() -> None:
 
     n = 100
     output = SIR_w_breaks(theta, n, seed=model_seed)
-
+    
     assert np.isclose(output, expected_output).all()

tests/test_examples/test_sir_python.py

@@ -34,9 +34,9 @@
 else:
     expected_params = np.array([[0.24, 0.26], [0.37, 0.21], [0.43, 0.14], [0.11, 0.04]])
 
-MAX_FLOAT32 = np.finfo(np.float32).max
-MIN_FLOAT32 = np.finfo(np.float32).min
-EPS_FLOAT32 = np.finfo(np.float32).eps
+MAX_FLOAT32: np.float64 = np.finfo(np.float32).max
+MIN_FLOAT32: np.float64 = np.finfo(np.float32).min
+EPS_FLOAT32: np.float64 = np.finfo(np.float32).eps
 
 
 def test_xgboost_2d() -> None: