added type annotations

src/pint/simulation.py

@@ -2,6 +2,8 @@
 """
 from collections import OrderedDict
 from copy import deepcopy
+from typing import Optional, Union, List, Dict
+import pathlib
 
 import astropy.units as u
 import numpy as np
@@ -22,7 +24,14 @@
 ]
 
 
-def zero_residuals(ts, model, *, subtract_mean=True, maxiter=10, tolerance=None):
+def zero_residuals(
+    ts: pint.toa.TOAs,
+    model: pint.models.timing_model.TimingModel,
+    *,
+    subtract_mean: bool = True,
+    maxiter: int = 10,
+    tolerance: Optional[u.Quantity] = None,
+):
     """Use a model to adjust a TOAs object, setting residuals to 0 iteratively.
 
     Parameters
@@ -62,7 +71,11 @@ def zero_residuals(ts, model, *, subtract_mean=True, maxiter=10, tolerance=None)
         )
 
 
-def get_fake_toa_clock_versions(model, include_bipm=False, include_gps=True):
+def get_fake_toa_clock_versions(
+    model: pint.models.timing_model.TimingModel,
+    include_bipm: bool = False,
+    include_gps: bool = True,
+) -> dict:
     """Get the clock settings (corrections, etc) for fake TOAs
 
     Parameters
@@ -75,6 +88,10 @@ def get_fake_toa_clock_versions(model, include_bipm=False, include_gps=True):
     include_gps : bool, optional
         Whether or not to disable UTC(GPS)->UTC clock correction
         (see :class:`pint.observatory.topo_obs.TopoObs`)
+
+    Returns
+    -------
+    dict
     """
     bipm_version = bipm_default
     if model["CLOCK"].value is not None:
@@ -109,13 +126,13 @@ def get_fake_toa_clock_versions(model, include_bipm=False, include_gps=True):
 
 
 def make_fake_toas(
-    ts,
-    model,
-    add_noise=False,
-    add_correlated_noise=False,
-    name="fake",
-    subtract_mean=True,
-):
+    ts: pint.toa.TOAs,
+    model: pint.models.timing_model.TimingModel,
+    add_noise: bool = False,
+    add_correlated_noise: bool = False,
+    name: str = "fake",
+    subtract_mean: bool = True,
+) -> pint.toa.TOAs:
     """Make toas from an array of times
 
     Can include alternating frequencies if fed an array of frequencies,
@@ -166,8 +183,23 @@ def make_fake_toas(
     return tsim
 
 
-def update_fake_dms(model, ts, dm_error, add_noise):
-    """Update simulated wideband DM information in TOAs."""
+def update_fake_dms(
+    model: pint.models.timing_model.TimingModel,
+    ts: pint.toa.TOAs,
+    dm_error: u.Quantity,
+    add_noise: bool,
+) -> pint.toa.TOAs:
+    """Update simulated wideband DM information in TOAs.
+
+    Parameters
+    ----------
+    model: pint.models.timing_model.TimingModel
+    ts : pint.toa.TOAs
+        Input TOAs
+    dm_error: u.Quantity
+    add_noise : bool, optional
+        Add noise to the DMs (otherwise `dm_error` just populates the column)
+    """
     toas = deepcopy(ts)
 
     dm_errors = dm_error * np.ones(len(toas))
@@ -187,25 +219,25 @@ def update_fake_dms(model, ts, dm_error, add_noise):
 
 
 def make_fake_toas_uniform(
-    startMJD,
-    endMJD,
-    ntoas,
-    model,
-    fuzz=0,
-    freq=1400 * u.MHz,
-    obs="GBT",
-    error=1 * u.us,
-    add_noise=False,
-    add_correlated_noise=False,
-    wideband=False,
-    wideband_dm_error=1e-4 * pint.dmu,
-    name="fake",
-    include_bipm=False,
-    include_gps=True,
-    multi_freqs_in_epoch=False,
-    flags=None,
-    subtract_mean=True,
-):
+    startMJD: float | u.Quantity | time.Time,
+    endMJD: float | u.Quantity | time.Time,
+    ntoas: int,
+    model: pint.models.timing_model.TimingModel,
+    fuzz: u.Quantity = 0,
+    freq: u.Quantity = 1400 * u.MHz,
+    obs: str = "GBT",
+    error: u.Quantity = 1 * u.us,
+    add_noise: bool = False,
+    add_correlated_noise: bool = False,
+    wideband: bool = False,
+    wideband_dm_error: u.Quantity = 1e-4 * pint.dmu,
+    name: str = "fake",
+    include_bipm: bool = False,
+    include_gps: bool = True,
+    multi_freqs_in_epoch: bool = False,
+    flags: Optional[dict] = None,
+    subtract_mean: bool = True,
+) -> pint.toa.TOAs:
     """Simulate uniformly spaced TOAs.
 
     Parameters
@@ -329,22 +361,22 @@ def make_fake_toas_uniform(
 
 
 def make_fake_toas_fromMJDs(
-    MJDs,
-    model,
-    freq=1400 * u.MHz,
-    obs="GBT",
-    error=1 * u.us,
-    add_noise=False,
-    add_correlated_noise=False,
-    wideband=False,
-    wideband_dm_error=1e-4 * pint.dmu,
-    name="fake",
-    include_bipm=False,
-    include_gps=True,
-    multi_freqs_in_epoch=False,
-    flags=None,
-    subtract_mean=True,
-):
+    MJDs: u.Quantity | time.Time | np.ndarray,
+    model: pint.models.timing_model.TimingModel,
+    freq: u.Quantity = 1400 * u.MHz,
+    obs: str = "GBT",
+    error: u.Quantity = 1 * u.us,
+    add_noise: bool = False,
+    add_correlated_noise: bool = False,
+    wideband: bool = False,
+    wideband_dm_error: u.Quantity = 1e-4 * pint.dmu,
+    name: str = "fake",
+    include_bipm: bool = False,
+    include_gps: bool = True,
+    multi_freqs_in_epoch: bool = False,
+    flags: Optional[dict] = None,
+    subtract_mean: bool = True,
+) -> pint.toa.TOAs:
     """Simulate TOAs from a list of MJDs
 
     Parameters
@@ -464,13 +496,13 @@ def make_fake_toas_fromMJDs(
 
 
 def make_fake_toas_fromtim(
-    timfile,
-    model,
-    add_noise=False,
-    add_correlated_noise=False,
-    name="fake",
-    subtract_mean=True,
-):
+    timfile: str | List[str] | pathlib.Path,
+    model: pint.models.timing_model.TimingModel,
+    add_noise: bool = False,
+    add_correlated_noise: bool = False,
+    name: str = "fake",
+    subtract_mean: bool = True,
+) -> pint.toa.TOAs:
     """Simulate fake TOAs with the same times as an input tim file
 
     Parameters
@@ -526,8 +558,13 @@ def make_fake_toas_fromtim(
 
 
 def calculate_random_models(
-    fitter, toas, Nmodels=100, keep_models=True, return_time=False, params="all"
-):
+    fitter: pint.fitter.Fitter,
+    toas: pint.toa.TOAs,
+    Nmodels: int = 100,
+    keep_models: bool = True,
+    return_time: bool = False,
+    params: str = "all",
+) -> (np.ndarray, Optional[list]):
     """
     Calculates random models based on the covariance matrix of the `fitter` object.
 
@@ -650,7 +687,13 @@ def calculate_random_models(
     return (dphase, random_models) if keep_models else dphase
 
 
-def _get_freqs_and_times(start, end, ntoas, freqs, multi_freqs_in_epoch=True):
+def _get_freqs_and_times(
+    start: float | u.Quantity | time.Time,
+    end: float | u.Quantity | time.Time,
+    ntoas: int,
+    freqs: u.Quantity,
+    multi_freqs_in_epoch: bool = True,
+) -> (float | u.Quantity | time.Time, np.ndarray):
     freqs = np.atleast_1d(freqs)
     assert (
         len(freqs.shape) == 1 and len(freqs) <= ntoas