diff --git a/src/pint/gridutils.py b/src/pint/gridutils.py index 517c1fa17..41df3eac8 100644 --- a/src/pint/gridutils.py +++ b/src/pint/gridutils.py @@ -1,10 +1,14 @@ """Tools for building chi-squared grids.""" + +from ast import Dict import concurrent.futures import copy import multiprocessing import subprocess import sys +from typing import List, Tuple, Optional, Dict + import numpy as np from loguru import logger as log @@ -15,17 +19,18 @@ tqdm = None from astropy.utils.console import ProgressBar +from astropy import units as u from pint import fitter __all__ = ["doonefit", "grid_chisq", "grid_chisq_derived"] -def hostinfo(): +def hostinfo() -> str: return subprocess.check_output("uname -a", shell=True) -def set_log(logger_): +def set_log(logger_: "loguru._logger.Logger") -> None: global log log = logger_ @@ -33,7 +38,7 @@ def set_log(logger_): class WrappedFitter: """Worker class to compute one fit with specified parameters fixed but passing other parameters to fit_toas()""" - def __init__(self, ftr, **fitargs): + def __init__(self, ftr: fitter.Fitter, **fitargs) -> None: """Worker class to computes one fit with specified parameters fixed Parameters @@ -45,7 +50,9 @@ def __init__(self, ftr, **fitargs): self.ftr = ftr self.fitargs = fitargs - def doonefit(self, parnames, parvalues, extraparnames=[]): + def doonefit( + self, parnames: List[str], parvalues: List[str], extraparnames: List[str] = [] + ) -> Tuple[float, List[float]]: """Worker process that computes one fit with specified parameters fixed Parameters @@ -105,7 +112,12 @@ def doonefit(self, parnames, parvalues, extraparnames=[]): return chi2, extraparvalues -def doonefit(ftr, parnames, parvalues, extraparnames=[]): +def doonefit( + ftr: fitter.Fitter, + parnames: List[str], + parvalues: List[u.Quantity], + extraparnames: List[str] = [], +) -> Tuple[float, List[float]]: """Worker process that computes one fit with specified parameters fixed Parameters @@ -153,16 +165,16 @@ def doonefit(ftr, parnames, parvalues, extraparnames=[]): def grid_chisq( - ftr, - parnames, - parvalues, - extraparnames=[], - executor=None, - ncpu=None, - chunksize=1, - printprogress=True, + ftr: fitter.Fitter, + parnames: List[str], + parvalues: List[u.Quantity], + extraparnames: List[str] = [], + executor: Optional[concurrent.futures.Executor] = None, + ncpu: Optional[int] = None, + chunksize: int = 1, + printprogress: bool = True, **fitargs, -): +) -> Tuple[np.ndarray, Dict[str, np.ndarray]]: """Compute chisq over a grid of parameters Parameters @@ -379,17 +391,17 @@ def grid_chisq( def grid_chisq_derived( - ftr, - parnames, - parfuncs, - gridvalues, - extraparnames=[], - executor=None, - ncpu=None, - chunksize=1, - printprogress=True, + ftr: fitter.Fitter, + parnames: List[str], + parfuncs: List[callable], + gridvalues: List[np.ndarray], + extraparnames: List[str] = [], + executor: Optional[concurrent.futures.Executor] = None, + ncpu: Optional[int] = None, + chunksize: int = 1, + printprogress: bool = True, **fitargs, -): +) -> Tuple[np.ndarray, List[np.ndarray], Dict[str, np.ndarray]]: """Compute chisq over a grid of derived parameters Parameters @@ -575,16 +587,16 @@ def grid_chisq_derived( def tuple_chisq( - ftr, - parnames, - parvalues, - extraparnames=[], - executor=None, - ncpu=None, - chunksize=1, - printprogress=True, + ftr: fitter.Fitter, + parnames: List[str], + parvalues: List[np.ndarray], + extraparnames: List[str] = [], + executor: Optional[concurrent.futures.Executor] = None, + ncpu: Optional[int] = None, + chunksize: int = 1, + printprogress: bool = True, **fitargs, -): +) -> Tuple[np.ndarray, Dict[str, np.ndarray]]: """Compute chisq over a list of parameter tuples Parameters @@ -760,17 +772,17 @@ def tuple_chisq( def tuple_chisq_derived( - ftr, - parnames, - parfuncs, - parvalues, - extraparnames=[], - executor=None, - ncpu=None, - chunksize=1, - printprogress=True, + ftr: fitter.Fitter, + parnames: List[str], + parfuncs: List[callable], + parvalues: List[np.ndarray], + extraparnames: List[str] = [], + executor: Optional[concurrent.futures.Executor] = None, + ncpu: Optional[int] = None, + chunksize: int = 1, + printprogress: bool = True, **fitargs, -): +) -> Tuple[np.ndarray, List, Dict[str, np.ndarray]]: """Compute chisq over a list of derived parameter tuples Parameters