Add type checking with ty

.github/workflows/codestyle.yml

@@ -40,3 +40,8 @@ jobs:
           set -euo pipefail
           uv run ruff check . --config ./pyproject.toml
           uv run ruff format --diff . --config ./pyproject.toml
+
+      - name: Check types with ty
+        run: |
+          set -euo pipefail
+          uv run ty check .

.pre-commit-config.yaml

@@ -22,3 +22,11 @@ repos:
     # Run the formatter.
     - id: ruff-format
       types_or: [ python, pyi, jupyter ]
+- repo: local
+  hooks:
+    - id: ty
+      name: ty type checker
+      entry: uv run ty check
+      language: system
+      types: [python]
+      pass_filenames: false

CHANGELOG.md

@@ -1,5 +1,7 @@
 # HEAD
 
+-   Add automatic type-checking with ty [#481](https://github.com/litebird/litebird_sim/pull/481)
+
 -   Stop using sphinxcontrib-asciinema [#478](https://github.com/litebird/litebird_sim/pull/478)
 
 -   TODs handling improved [#472](https://github.com/litebird/litebird_sim/pull/472)

docs/source/profiling.rst

@@ -59,7 +59,7 @@ Here is an example::
     from litebird_sim import TimeProfiler, profile_list_to_speedscope
     from time import sleep
 
-    perf_list = []  # type: list[TimeProfiler]
+    perf_list: list[TimeProfiler] = []
 
     # First code block to profile
     with TimeProfiler(name="function_A") as perf:

litebird_sim/__init__.py

@@ -203,7 +203,7 @@
 
 # Privilege TBB over OpenPM and the internal Numba implementation of a
 # work queue
-numba.config.THREADING_LAYER_CONFIG = ["tbb", "omp", "workqueue"]
+numba.config.THREADING_LAYER_CONFIG = ["tbb", "omp", "workqueue"]  # type: ignore
 
 
 __all__ = [

litebird_sim/bandpass_template_module.py

@@ -106,7 +106,7 @@ def beam_throughtput(freqs):
 
 # Define bandpass profile
 def bandpass_profile(
-    freqs: np.ndarray | None = None,
+    freqs: np.ndarray,
     bandpass: dict | None = None,
     include_beam_throughput: bool | None = None,
 ):
@@ -132,6 +132,8 @@ def bandpass_profile(
 
     profile = np.ones_like(freqs)
 
+    assert bandpass is not None, "bandpass dict is required"
+
     if "bandpass_file" in bandpass.keys():
         try:
             f, profile = np.loadtxt(

litebird_sim/bandpasses.py

@@ -1,4 +1,4 @@
-from dataclasses import dataclass
+from dataclasses import dataclass, fields
 
 import numpy as np
 import scipy as sp
@@ -84,13 +84,13 @@ class BandPassInfo:
     name: str = ""
     bandtype: str = "top-hat"
     normalize: bool = False
-    model: str = None
+    model: str | None = None
 
     alpha_exp: float = 1
     beta_exp: float = 1
     cosine_apo_length: float = 5
     cheby_poly_order: int = 3
-    cheby_ripple_dB: int = 3
+    cheby_ripple_dB: int | float = 3
 
     def __post_init__(self):
         """
@@ -240,6 +240,12 @@ def find_central_frequency(self):
                 / self.get_normalization()
             )
 
+    @classmethod
+    def from_dict(cls, d: dict) -> "BandPassInfo":
+        """Create a BandPassInfo object from a dictionary."""
+        field_names = {f.name for f in fields(cls)}
+        return cls(**{k: v for k, v in d.items() if k in field_names})
+
     @staticmethod
     def from_imo(imo: Imo, url: UUID | str):
         """Create a `BandPassInfo` object from a definition in the IMO

litebird_sim/beam_convolution.py

@@ -106,6 +106,7 @@ def add_convolved_sky_to_one_detector(
         beam_mmax = beam_alms_det.mmax
 
         default_lmax = min(sky_lmax, beam_lmax)
+        assert beam_mmax is not None
         default_mmax = min(default_lmax - 4, beam_mmax)
 
         logging.warning(
@@ -303,12 +304,16 @@ def add_convolved_sky(
         if input_sky_names is None:
             sky_alms_det = sky_alms
         else:
+            assert isinstance(sky_alms, dict)
             sky_alms_det = sky_alms[input_sky_names[detector_idx]]
+        assert isinstance(sky_alms_det, SphericalHarmonics)
 
         if input_beam_names is None:
             beam_alms_det = beam_alms
         else:
+            assert isinstance(beam_alms, dict)
             beam_alms_det = beam_alms[input_beam_names[detector_idx]]
+        assert isinstance(beam_alms_det, SphericalHarmonics)
 
         if mueller_hwp is None:
             mueller_matrix = None
@@ -331,12 +336,12 @@ def add_convolved_sky(
 def add_convolved_sky_to_observations(
     observations: Observation | list[Observation],
     sky_alms: (
-        SphericalHarmonics | dict[str, SphericalHarmonics]
+        SphericalHarmonics | dict[str, SphericalHarmonics] | None
     ),  # at some point optional, taken from the obs
     beam_alms: (
-        SphericalHarmonics | dict[str, SphericalHarmonics]
+        SphericalHarmonics | dict[str, SphericalHarmonics] | None
     ),  # at some point optional, taken from the obs
-    pointings: npt.ArrayLike | list[npt.ArrayLike] | None = None,
+    pointings: npt.NDArray | list[npt.NDArray] | None = None,
     hwp: HWP | None = None,
     input_sky_alms_in_galactic: bool = True,
     convolution_params: BeamConvolutionParameters | None = None,
@@ -396,20 +401,22 @@ def add_convolved_sky_to_observations(
 
     if sky_alms is None:
         try:
-            sky_alms = observations[0].sky
+            assert obs_list, "No observations available"
+            sky_alms = obs_list[0].sky
         except AttributeError:
             msg = "'sky_alms' is None and nothing is found in the observation. You should either pass the spherical harmonics, or store them in the observations if 'mbs' is used."
             raise AttributeError(msg)
-        assert sky_alms["type"] == "alms", (
+        assert isinstance(sky_alms, dict) and sky_alms["type"] == "alms", (
             "'sky_alms' should be of type 'alms'. Use 'store_alms' of 'MbsParameters' to make it so."
         )
 
     if beam_alms is None:
         try:
-            beam_alms = observations[0].blms
+            beam_alms = obs_list[0].blms
         except AttributeError:
             msg = "'beam_alms' is None and nothing is found in the observation. You should either pass the spherical harmonics of the beam, or store them in the observations if 'get_gauss_beam_alms' is used."
             raise AttributeError(msg)
+    assert beam_alms is not None
 
     for cur_obs, cur_ptg in zip(obs_list, ptg_list):
         # Determine input sky names

litebird_sim/beam_synthesis.py

@@ -65,10 +65,10 @@ def gauss_beam_to_alm(
     lmax: int,
     mmax: int,
     fwhm_rad: float,
-    ellipticity: float | None = 1.0,
-    psi_ell_rad: float | None = 0.0,
-    psi_pol_rad: float | None = 0.0,
-    cross_polar_leakage: float | None = 0.0,
+    ellipticity: float = 1.0,
+    psi_ell_rad: float = 0.0,
+    psi_pol_rad: float = 0.0,
+    cross_polar_leakage: float = 0.0,
 ) -> SphericalHarmonics:
     """
     Compute spherical harmonics coefficients a_ℓm representing a Gaussian beam.
@@ -238,6 +238,7 @@ def generate_gauss_beam_alms(
 
     blms = {}
 
+    assert observation.name is not None
     if channels is None:
         # Use detectors from observations
         for detector_idx, det_name in enumerate(observation.name):

litebird_sim/detectors.py

@@ -188,14 +188,15 @@ class DetectorInfo:
     quat: Any = None
     pol_angle_rad: float = 0.0
     pol_efficiency: float = 1.0
-    mueller_hwp: None | np.ndarray | dict[np.ndarray] = None
-    mueller_hwp_solver: None | np.ndarray | dict[np.ndarray] = None
-    jones_hwp: None | np.ndarray | dict[np.ndarray] = None
-    jones_hwp_solver: None | np.ndarray | dict[np.ndarray] = None
+    mueller_hwp: None | np.ndarray | dict[str, np.ndarray] = None
+    mueller_hwp_solver: None | np.ndarray | dict[str, np.ndarray] = None
+    jones_hwp: None | np.ndarray | dict[str, np.ndarray] = None
+    jones_hwp_solver: None | np.ndarray | dict[str, np.ndarray] = None
     pointing_theta_phi_psi_deg: None | np.ndarray = None
     pointing_u_v: None | np.ndarray = None
     g_one_over_k: float = 0.0
     amplitude_2f_k: float = 0.0
+    band: BandPassInfo | None = None
 
     def __post_init__(self):
         if self.quat is None:
@@ -204,6 +205,7 @@ def __post_init__(self):
         self.quat = normalize_time_dependent_quaternion(self.quat)
 
         if isinstance(self.band_freqs_ghz, np.ndarray):
+            assert self.band_weights is not None
             assert len(self.band_freqs_ghz) == len(self.band_weights)
 
         # Warn if mueller_hwp is not a 4x4 numpy array
@@ -341,6 +343,7 @@ class FreqChannelInfo:
     number_of_detectors: int = 0
     detector_names: list[str] = field(default_factory=list)
     detector_objs: list[UUID] = field(default_factory=list)
+    band: BandPassInfo | None = None
 
     """A data class representing the configuration of a frequency channel in LiteBIRD.
 

litebird_sim/dipole.py

@@ -137,31 +137,31 @@ def add_dipole_for_one_detector(
     dipole_type: DipoleType,
 ):
     if dipole_type == DipoleType.LINEAR:
-        for i in prange(len(tod_det)):
+        for i in prange(len(tod_det)):  # type: ignore[not-iterable]
             tod_det[i] += compute_dipole_for_one_sample_linear(
                 theta=theta_phi_det[i, 0],
                 phi=theta_phi_det[i, 1],
                 v_km_s=velocity[i],
                 t_cmb_k=t_cmb_k,
             )
     elif dipole_type == DipoleType.QUADRATIC_EXACT:
-        for i in prange(len(tod_det)):
+        for i in prange(len(tod_det)):  # type: ignore[not-iterable]
             tod_det[i] += compute_dipole_for_one_sample_quadratic_exact(
                 theta=theta_phi_det[i, 0],
                 phi=theta_phi_det[i, 1],
                 v_km_s=velocity[i],
                 t_cmb_k=t_cmb_k,
             )
     elif dipole_type == DipoleType.TOTAL_EXACT:
-        for i in prange(len(tod_det)):
+        for i in prange(len(tod_det)):  # type: ignore[not-iterable]
             tod_det[i] += compute_dipole_for_one_sample_total_exact(
                 theta=theta_phi_det[i, 0],
                 phi=theta_phi_det[i, 1],
                 v_km_s=velocity[i],
                 t_cmb_k=t_cmb_k,
             )
     elif dipole_type == DipoleType.QUADRATIC_FROM_LIN_T:
-        for i in prange(len(tod_det)):
+        for i in prange(len(tod_det)):  # type: ignore[not-iterable]
             tod_det[i] += compute_dipole_for_one_sample_quadratic_from_lin_t(
                 theta=theta_phi_det[i, 0],
                 phi=theta_phi_det[i, 1],
@@ -171,7 +171,7 @@ def add_dipole_for_one_detector(
             )
     elif dipole_type == DipoleType.TOTAL_FROM_LIN_T:
         planck_t0 = planck(nu_hz, t_cmb_k)
-        for i in prange(len(tod_det)):
+        for i in prange(len(tod_det)):  # type: ignore[not-iterable]
             tod_det[i] += compute_dipole_for_one_sample_total_from_lin_t(
                 theta=theta_phi_det[i, 0],
                 phi=theta_phi_det[i, 1],

litebird_sim/distribute.py

@@ -226,7 +226,7 @@ def distribute_optimally(elements, num_of_groups, weight_fn=None) -> list[Span]:
 
     max_weight = _partition(elements, len(elements), num_of_groups, weight_fn)
 
-    result = []  # type: list[Span]
+    result: list[Span] = []
     start_idx = 0
     weight = 0
     cur_num = 0

litebird_sim/gaindrifts.py

@@ -1,7 +1,8 @@
-import numpy as np
-
-from enum import IntEnum
+import logging as log
 from dataclasses import dataclass
+from enum import IntEnum
+
+import numpy as np
 
 from .observations import Observation
 from .seeding import regenerate_or_check_detector_generators
@@ -22,7 +23,7 @@ class GainDriftType(IntEnum):
 
     LINEAR_GAIN = 0
     THERMAL_GAIN = 1
-    # SLOW_GAIN = 2 # Remains to be implemented
+    SLOW_GAIN = 2  # Remains to be implemented
 
 
 class SamplingDist(IntEnum):
@@ -202,8 +203,8 @@ def _get_psd(
 def _noise_timestream(
     tod_size: int,
     sampling_freq_hz: float,
-    drift_params: GainDriftParams = None,
-    random: np.random.Generator = None,
+    drift_params: GainDriftParams | None = None,
+    random: np.random.Generator | None = None,
 ) -> np.ndarray:
     """The function to generate the thermal noise time stream with
     :math:`1/f` power spectral density.
@@ -279,9 +280,9 @@ def _noise_timestream(
 def apply_gaindrift_for_one_detector(
     det_tod: np.ndarray,
     sampling_freq_hz: float,
-    drift_params: GainDriftParams = None,
-    noise_timestream: np.ndarray = None,
-    random: np.random.Generator = None,
+    drift_params: GainDriftParams | None = None,
+    noise_timestream: np.ndarray | None = None,
+    random: np.random.Generator | None = None,
 ):
     """This function applies the gain drift on the TOD corresponding to only one
     detector.
@@ -389,7 +390,7 @@ def apply_gaindrift_for_one_detector(
         det_tod *= _responsivity_function(dT)
 
     elif drift_params.drift_type == GainDriftType.SLOW_GAIN:
-        # !!! Remains to be implemented
+        log.warning("Slow gain drift is not yet implemented.")
         pass
     else:
         raise ValueError(
@@ -401,9 +402,9 @@ def apply_gaindrift_for_one_detector(
 def apply_gaindrift_to_tod(
     tod: np.ndarray,
     sampling_freq_hz: float,
-    drift_params: GainDriftParams = None,
-    focalplane_attr: list | np.ndarray = None,
-    dets_random: np.random.Generator | None = None,
+    drift_params: GainDriftParams | None = None,
+    focalplane_attr: list | np.ndarray | None = None,
+    dets_random: list[np.random.Generator] | None = None,
 ):
     """The function to apply the gain drift to all the detectors of a given TOD object.
 
@@ -448,6 +449,8 @@ def apply_gaindrift_to_tod(
     if drift_params is None:
         drift_params = GainDriftParams()
 
+    assert dets_random is not None, "dets_random is required"
+
     tod_size = len(tod[0])
 
     if drift_params.drift_type == GainDriftType.LINEAR_GAIN:
@@ -501,10 +504,10 @@ def apply_gaindrift_to_tod(
 
 def apply_gaindrift_to_observations(
     observations: Observation | list[Observation],
-    drift_params: GainDriftParams = None,
+    drift_params: GainDriftParams | None = None,
     user_seed: int | None = None,
     component: str = "tod",
-    dets_random: list[np.random.Generator] = None,
+    dets_random: list[np.random.Generator] | None = None,
 ):
     """
     Apply gain drift to one or more observations.