diff --git a/CHANGELOG-unreleased.md b/CHANGELOG-unreleased.md index 2a705d81a..058351dc8 100644 --- a/CHANGELOG-unreleased.md +++ b/CHANGELOG-unreleased.md @@ -14,7 +14,8 @@ the released changes. - Changed default value of `FDJUMPLOG` to `Y` - Bumped `black` version to 24.x ### Added -- Type hints in `pint.derived_quantities` and `pint.residuals` +- Type hints in `pint.derived_quantities`, `pint.modelutils`, `pint.binaryconvert`, `pint.config`, +`pint.erfautils`, `pint.fits_utils`, `pint.logging` and `pint.residuals` - Doing `model.par = something` will try to assign to `par.quantity` or `par.value` but will give warning - `PLChromNoise` component to model chromatic red noise with a power law spectrum ### Fixed diff --git a/src/pint/binaryconvert.py b/src/pint/binaryconvert.py index 9edd2f697..3410cdf53 100644 --- a/src/pint/binaryconvert.py +++ b/src/pint/binaryconvert.py @@ -12,12 +12,14 @@ import copy from uncertainties import ufloat, umath from loguru import logger as log +from typing import Optional, List, Tuple, Union from pint import Tsun from pint.models.binary_bt import BinaryBT from pint.models.binary_dd import BinaryDD, BinaryDDS, BinaryDDH from pint.models.binary_ddk import BinaryDDK from pint.models.binary_ell1 import BinaryELL1, BinaryELL1H, BinaryELL1k +import pint.models # output types # DDGR is not included as there is not a well-defined way to get a unique output @@ -27,7 +29,7 @@ __all__ = ["convert_binary"] -def _M2SINI_to_orthometric(model): +def _M2SINI_to_orthometric(model: pint.models.TimingModel) -> Tuple[u.Quantity]: """Convert from standard Shapiro delay (M2, SINI) to orthometric (H3, H4, STIGMA) Uses Eqns. 12, 20, 21 from Freire and Wex (2010) @@ -76,7 +78,7 @@ def _M2SINI_to_orthometric(model): return stigma.n, h3.n * u.s, h4.n * u.s, stigma_unc, h3_unc, h4_unc -def _orthometric_to_M2SINI(model): +def _orthometric_to_M2SINI(model: pint.models.TimingModel) -> Tuple[u.Quantity]: """Convert from orthometric (H3, H4, STIGMA) to standard Shapiro delay (M2, SINI) Inverts Eqns. 12, 20, 21 from Freire and Wex (2010) @@ -136,7 +138,7 @@ def _orthometric_to_M2SINI(model): return m2.n * u.Msun, sini.n, m2_unc, sini_unc -def _SINI_to_SHAPMAX(model): +def _SINI_to_SHAPMAX(model: pint.models.TimingModel) -> Tuple[u.Quantity]: """Convert from standard SINI to alternate SHAPMAX parameterization Also propagates uncertainties if present @@ -158,7 +160,7 @@ def _SINI_to_SHAPMAX(model): return shapmax.n, shapmax.s if shapmax.s > 0 else None -def _SHAPMAX_to_SINI(model): +def _SHAPMAX_to_SINI(model: pint.models.TimingModel) -> Tuple[u.Quantity]: """Convert from alternate SHAPMAX to SINI parameterization Also propagates uncertainties if present @@ -180,7 +182,7 @@ def _SHAPMAX_to_SINI(model): return sini.n, sini.s if sini.s > 0 else None -def _from_ELL1(model): +def _from_ELL1(model: pint.models.TimingModel) -> Tuple[u.Quantity]: """Convert from ELL1 parameterization to standard orbital parameterization Converts using Eqns. 1, 2, and 3 from Lange et al. (2001) @@ -270,7 +272,7 @@ def _from_ELL1(model): ) -def _to_ELL1(model): +def _to_ELL1(model: pint.models.TimingModel) -> Tuple[u.Quantity]: """Convert from standard orbital parameterization to ELL1 parameterization Converts using Eqns. 1, 2, and 3 from Lange et al. (2001) @@ -353,7 +355,7 @@ def _to_ELL1(model): ) -def _ELL1_to_ELL1k(model): +def _ELL1_to_ELL1k(model: pint.models.TimingModel) -> Tuple[u.Quantity]: """Convert from ELL1 EPS1DOT/EPS2DOT to ELL1k LNEDOT/OMDOT Parameters @@ -391,7 +393,7 @@ def _ELL1_to_ELL1k(model): return lnedot.n / u.s, (omdot.n / u.s).to(u.deg / u.yr), lnedot_unc, omdot_unc -def _ELL1k_to_ELL1(model): +def _ELL1k_to_ELL1(model: pint.models.TimingModel) -> Tuple[u.Quantity]: """Convert from ELL1k LNEDOT/OMDOT to ELL1 EPS1DOT/EPS2DOT Parameters @@ -428,7 +430,7 @@ def _ELL1k_to_ELL1(model): return eps1dot.n / u.s, eps2dot.n / u.s, eps1dot_unc, eps2dot_unc -def _DDGR_to_PK(model): +def _DDGR_to_PK(model: pint.models.TimingModel) -> Tuple[u.Quantity]: """Convert DDGR model to equivalent PK parameters Uses ``uncertainties`` module to propagate uncertainties @@ -507,7 +509,11 @@ def _DDGR_to_PK(model): return pbdot, gamma, omegadot, s, r, Dr, Dth -def _transfer_params(inmodel, outmodel, badlist=[]): +def _transfer_params( + inmodel: pint.models.TimingModel, + outmodel: pint.models.TimingModel, + badlist: list = [], +) -> None: """Transfer parameters between an input and output model, excluding certain parameters Parameters (input or output) that are :class:`~pint.models.parameter.funcParameter` are not copied @@ -535,7 +541,13 @@ def _transfer_params(inmodel, outmodel, badlist=[]): ) -def convert_binary(model, output, NHARMS=3, useSTIGMA=False, KOM=0 * u.deg): +def convert_binary( + model: pint.models.TimingModel, + output: str, + NHARMS: int = 3, + useSTIGMA: bool = False, + KOM: u.Quantity = 0 * u.deg, +) -> pint.models.TimingModel: """ Convert between binary models diff --git a/src/pint/config.py b/src/pint/config.py index 09754b852..a24bec249 100644 --- a/src/pint/config.py +++ b/src/pint/config.py @@ -7,7 +7,7 @@ # location of tim, par files installed via pkg_resources -def datadir(): +def datadir() -> str: """Location of the PINT data (par and tim files) Returns @@ -18,7 +18,7 @@ def datadir(): return pkg_resources.resource_filename(__name__, "data/") -def examplefile(filename): +def examplefile(filename: str) -> str: """Full path to the requested PINT example data file Parameters @@ -40,7 +40,7 @@ def examplefile(filename): ) -def runtimefile(filename): +def runtimefile(filename: str) -> str: """Full path to the requested PINT runtime (clock etc) data file Parameters diff --git a/src/pint/erfautils.py b/src/pint/erfautils.py index eb42f532d..01270a39b 100644 --- a/src/pint/erfautils.py +++ b/src/pint/erfautils.py @@ -4,11 +4,11 @@ import numpy as np from astropy import table +from astropy.coordinates import EarthLocation from pint.pulsar_mjd import Time from pint.utils import PosVel - __all__ = ["gcrs_posvel_from_itrf"] SECS_PER_DAY = erfa.DAYSEC @@ -24,7 +24,9 @@ asec2rad = 4.84813681109536e-06 -def gcrs_posvel_from_itrf(loc, toas, obsname="obs"): +def gcrs_posvel_from_itrf( + loc: EarthLocation, toas: "pint.toa.TOAs", obsname: str = "obs" +) -> PosVel: """Return a list of PosVel instances for the observatory at the TOA times. Observatory location should be given in the loc argument as an astropy @@ -45,6 +47,17 @@ def gcrs_posvel_from_itrf(loc, toas, obsname="obs"): This version uses astropy's internal routines, which use IERS A data rather than the final IERS B values. These do differ, and yield results that are different by ~20 m. + + Parameters + ---------- + loc : astropy.coordinates.EarthLocation + toas : pint.toa.TOAs + obsname : str + + Returns + ------- + `pint.utils.PosVel` + """ unpack = False # If the input is a single TOA (i.e. a row from the table), diff --git a/src/pint/fits_utils.py b/src/pint/fits_utils.py index 61ab27d40..cb2f55621 100644 --- a/src/pint/fits_utils.py +++ b/src/pint/fits_utils.py @@ -5,21 +5,30 @@ from erfa import DAYSEC as SECS_PER_DAY +from astropy.io import fits + from pint.pulsar_mjd import fortran_float __all__ = ["read_fits_event_mjds", "read_fits_event_mjds_tuples"] -def read_fits_event_mjds_tuples(event_hdu, timecolumn="TIME"): +def read_fits_event_mjds_tuples( + event_hdu: fits.hdu.BinTableHDU, timecolumn: str = "TIME" +) -> np.ndarray: """Read a set of MJDs from a FITS HDU, with proper conversion of times to MJD The FITS time format is defined here: https://heasarc.gsfc.nasa.gov/docs/journal/timing3.html + Parameters + ---------- + event_hdu : fits.hdu.BinTableHDU + timecolumn : str, optional + Returns ------- - mjds: MJDs returned are tuples of two doubles (jd1, jd2), as use by + mjds: MJDs returned are tuples of two doubles (jd1, jd2), as used by astropy Time() objects. """ @@ -63,13 +72,24 @@ def read_fits_event_mjds_tuples(event_hdu, timecolumn="TIME"): ) -def read_fits_event_mjds(event_hdu, timecolumn="TIME"): +def read_fits_event_mjds( + event_hdu: fits.hdu.BinTableHDU, timecolumn: str = "TIME" +) -> np.ndarray: """Read a set of MJDs from a FITS HDU, with proper conversion of times to MJD The FITS time format is defined here: https://heasarc.gsfc.nasa.gov/docs/journal/timing3.html MJDs returned are double precision floats + + Parameters + ---------- + event_hdu : fits.hdu.BinTableHDU + timecolumn : str, optional + + Returns + ------- + mjds: np.ndarray """ event_hdr = event_hdu.header diff --git a/src/pint/gridutils.py b/src/pint/gridutils.py index 0df17ae0e..fce622f83 100644 --- a/src/pint/gridutils.py +++ b/src/pint/gridutils.py @@ -6,6 +6,8 @@ import subprocess import sys +from typing import List, Tuple, Optional, Dict + import numpy as np from loguru import logger as log @@ -16,17 +18,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_ @@ -34,7 +37,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 @@ -46,7 +49,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 @@ -106,7 +111,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 @@ -154,16 +164,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 @@ -380,17 +390,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 @@ -576,16 +586,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 @@ -761,17 +771,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 diff --git a/src/pint/logging.py b/src/pint/logging.py index e87a039bd..3e908f6dd 100644 --- a/src/pint/logging.py +++ b/src/pint/logging.py @@ -54,7 +54,8 @@ import sys import warnings from loguru import logger as log - +from typing import Union, Optional, Dict, List +import pint.types __all__ = ["LogFilter", "setup", "format", "levels", "get_level"] @@ -79,7 +80,14 @@ levels = ["TRACE", "DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"] -def showwarning(message, category, filename, lineno, file=None, line=None): +def showwarning( + message: Union[str, Warning], + category: Warning, + filename: str, + lineno: int, + file=None, + line=None, +) -> None: """ Function to allow ``loguru`` to capture warnings emitted by :func:`warnings.warn`. @@ -127,8 +135,13 @@ class LogFilter: Others that will only be seen once. Filtering of those is done on the basis of regular expressions. """ - def __init__(self, onlyonce=None, never=None, onlyonce_level="INFO"): - """ + def __init__( + self, + onlyonce: Optional[List[str]] = None, + never: Optional[List[str]] = None, + onlyonce_level: str = "INFO", + ) -> None: + r""" Define regexs for messages that will only be seen once. Use ``\S+`` for a variable that might change. If a message comes through with a new value for that variable, it will be seen. @@ -194,7 +207,7 @@ def __init__(self, onlyonce=None, never=None, onlyonce_level="INFO"): self.onlyonce_level = onlyonce_level - def filter(self, record): + def filter(self, record: Dict[str, str]) -> bool: """Filter the record based on ``record["message"]`` and ``record["level"]`` If this returns s,``False``, the message is not seen @@ -234,15 +247,15 @@ def __call__(self, record): def setup( - level="INFO", - sink=sys.stderr, - format=format, - filter=LogFilter(), - usecolors=True, - colors={"DEBUG": debug_color}, - capturewarnings=True, - removeprior=True, -): + level: str = "INFO", + sink: pint.types.file_like = sys.stderr, + format: str = format, + filter: callable = LogFilter(), + usecolors: bool = True, + colors: Dict[str, str] = {"DEBUG": debug_color}, + capturewarnings: bool = True, + removeprior: bool = True, +) -> int: """ Setup the PINT logging using ``loguru`` @@ -319,7 +332,7 @@ def setup( return loghandler -def get_level(starting_level_name, verbosity, quietness): +def get_level(starting_level_name: str, verbosity: int, quietness: int) -> str: """Get appropriate logging level given command-line input Parameters diff --git a/src/pint/modelutils.py b/src/pint/modelutils.py index 0a2842788..14cead276 100644 --- a/src/pint/modelutils.py +++ b/src/pint/modelutils.py @@ -1,11 +1,13 @@ from loguru import logger as log from pint.models.astrometry import AstrometryEquatorial, AstrometryEcliptic - +import pint.models # FIXME: shouldn't this be in the AstrometryEquatorial and AstrometryEcliptic classes? -def model_ecliptic_to_equatorial(model, force=False): +def model_ecliptic_to_equatorial( + model: pint.models.TimingModel, force: bool = False +) -> pint.models.TimingModel: r"""Converts Astrometry model component, Ecliptic to Equatorial Parameters @@ -55,7 +57,9 @@ def model_ecliptic_to_equatorial(model, force=False): return model -def model_equatorial_to_ecliptic(model, force=False): +def model_equatorial_to_ecliptic( + model: pint.models.TimingModel, force: bool = False +) -> pint.models.TimingModel: """Converts Astrometry model component, Equatorial to Ecliptic Parameters @@ -67,7 +71,7 @@ def model_equatorial_to_ecliptic(model, force=False): Returns ------- - model + model : `pint.models.TimingModel` new model with AstrometryEcliptic component """