diff --git a/.flake8 b/.flake8 index ae28ee95..07516cb3 100644 --- a/.flake8 +++ b/.flake8 @@ -1,4 +1,4 @@ [flake8] -ignore = E127,E201,E202,E203,E231,E252,E266,E402,E999,F841,W503,W605 +ignore = E127,E201,E202,E203,E231,E252,E266,E402,E999,F841,W503,W605,E704 max-line-length = 80 exclude = .git,docs,docsrc,scripts,cmdstanpy_tutorial.py,rtd_change_default_version.py diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 2c42518c..d5b1459b 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -44,7 +44,7 @@ jobs: strategy: matrix: os: [ubuntu-latest, macos-latest, windows-latest] - python-version: ["3.9", "3.10", "3.11", "3.12", "3.13", "3.14"] + python-version: ["3.10", "3.11", "3.12", "3.13", "3.14"] env: GITHUB_PAT: ${{ secrets.GITHUB_TOKEN }} steps: diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 7a5baefe..a1e3bbed 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -2,25 +2,25 @@ exclude: 'docsrc' fail_fast: true repos: - repo: https://github.com/pre-commit/pre-commit-hooks - rev: v3.2.0 + rev: v6.0.0 hooks: - id: check-yaml # isort should run before black as black sometimes tweaks the isort output - repo: https://github.com/PyCQA/isort - rev: 5.12.0 + rev: 7.0.0 hooks: - id: isort # https://github.com/python/black#version-control-integration - repo: https://github.com/psf/black - rev: 23.7.0 + rev: 25.9.0 hooks: - id: black - repo: https://github.com/pycqa/flake8 - rev: 3.9.2 + rev: 7.3.0 hooks: - id: flake8 - repo: https://github.com/pre-commit/mirrors-mypy - rev: v1.5.0 + rev: v1.18.2 hooks: - id: mypy additional_dependencies: [ numpy >= 1.22] diff --git a/cmdstanpy/cmdstan_args.py b/cmdstanpy/cmdstan_args.py index 9a7b2049..a37ebc17 100644 --- a/cmdstanpy/cmdstan_args.py +++ b/cmdstanpy/cmdstan_args.py @@ -5,14 +5,14 @@ import os from enum import Enum, auto from time import time -from typing import Any, Mapping, Optional, Union +from typing import Any, Mapping import numpy as np from numpy.random import default_rng from cmdstanpy.utils import cmdstan_path, cmdstan_version_before, get_logger -OptionalPath = Union[str, os.PathLike, None] +OptionalPath = str | os.PathLike | None class Method(Enum): @@ -52,19 +52,19 @@ class SamplerArgs: def __init__( self, - iter_warmup: Optional[int] = None, - iter_sampling: Optional[int] = None, + iter_warmup: int | None = None, + iter_sampling: int | None = None, save_warmup: bool = False, - thin: Optional[int] = None, - max_treedepth: Optional[int] = None, - metric_type: Optional[str] = None, - metric_file: Union[str, list[str], None] = None, - step_size: Union[float, list[float], None] = None, + thin: int | None = None, + max_treedepth: int | None = None, + metric_type: str | None = None, + metric_file: str | list[str] | None = None, + step_size: float | list[float] | None = None, adapt_engaged: bool = True, - adapt_delta: Optional[float] = None, - adapt_init_phase: Optional[int] = None, - adapt_metric_window: Optional[int] = None, - adapt_step_size: Optional[int] = None, + adapt_delta: float | None = None, + adapt_init_phase: int | None = None, + adapt_metric_window: int | None = None, + adapt_step_size: int | None = None, fixed_param: bool = False, num_chains: int = 1, ) -> None: @@ -74,8 +74,8 @@ def __init__( self.save_warmup = save_warmup self.thin = thin self.max_treedepth = max_treedepth - self.metric_type: Optional[str] = metric_type - self.metric_file: Union[str, list[str], None] = metric_file + self.metric_type: str | None = metric_type + self.metric_file: str | list[str] | None = metric_file self.step_size = step_size self.adapt_engaged = adapt_engaged self.adapt_delta = adapt_delta @@ -86,7 +86,7 @@ def __init__( self.diagnostic_file = None self.num_chains = num_chains - def validate(self, chains: Optional[int]) -> None: + def validate(self, chains: int | None) -> None: """ Check arguments correctness and consistency. @@ -295,16 +295,16 @@ class OptimizeArgs: def __init__( self, - algorithm: Optional[str] = None, - init_alpha: Optional[float] = None, - iter: Optional[int] = None, + algorithm: str | None = None, + init_alpha: float | None = None, + iter: int | None = None, save_iterations: bool = False, - tol_obj: Optional[float] = None, - tol_rel_obj: Optional[float] = None, - tol_grad: Optional[float] = None, - tol_rel_grad: Optional[float] = None, - tol_param: Optional[float] = None, - history_size: Optional[int] = None, + tol_obj: float | None = None, + tol_rel_obj: float | None = None, + tol_grad: float | None = None, + tol_rel_grad: float | None = None, + tol_param: float | None = None, + history_size: int | None = None, jacobian: bool = False, ) -> None: self.algorithm = algorithm or "" @@ -319,7 +319,7 @@ def __init__( self.history_size = history_size self.jacobian = jacobian - def validate(self, _chains: Optional[int] = None) -> None: + def validate(self, _chains: int | None = None) -> None: """ Check arguments correctness and consistency. """ @@ -383,13 +383,13 @@ class LaplaceArgs: """Arguments needed for laplace method.""" def __init__( - self, mode: str, draws: Optional[int] = None, jacobian: bool = True + self, mode: str, draws: int | None = None, jacobian: bool = True ) -> None: self.mode = mode self.jacobian = jacobian self.draws = draws - def validate(self, _chains: Optional[int] = None) -> None: + def validate(self, _chains: int | None = None) -> None: """Check arguments correctness and consistency.""" if not os.path.exists(self.mode): raise ValueError(f'Invalid path for mode file: {self.mode}') @@ -411,18 +411,18 @@ class PathfinderArgs: def __init__( self, - init_alpha: Optional[float] = None, - tol_obj: Optional[float] = None, - tol_rel_obj: Optional[float] = None, - tol_grad: Optional[float] = None, - tol_rel_grad: Optional[float] = None, - tol_param: Optional[float] = None, - history_size: Optional[int] = None, - num_psis_draws: Optional[int] = None, - num_paths: Optional[int] = None, - max_lbfgs_iters: Optional[int] = None, - num_draws: Optional[int] = None, - num_elbo_draws: Optional[int] = None, + init_alpha: float | None = None, + tol_obj: float | None = None, + tol_rel_obj: float | None = None, + tol_grad: float | None = None, + tol_rel_grad: float | None = None, + tol_param: float | None = None, + history_size: int | None = None, + num_psis_draws: int | None = None, + num_paths: int | None = None, + max_lbfgs_iters: int | None = None, + num_draws: int | None = None, + num_elbo_draws: int | None = None, save_single_paths: bool = False, psis_resample: bool = True, calculate_lp: bool = True, @@ -445,7 +445,7 @@ def __init__( self.psis_resample = psis_resample self.calculate_lp = calculate_lp - def validate(self, _chains: Optional[int] = None) -> None: + def validate(self, _chains: int | None = None) -> None: """ Check arguments correctness and consistency. """ @@ -514,7 +514,7 @@ def __init__(self, csv_files: list[str]) -> None: def validate( self, - chains: Optional[int] = None, # pylint: disable=unused-argument + chains: int | None = None, # pylint: disable=unused-argument ) -> None: """ Check arguments correctness and consistency. @@ -543,16 +543,16 @@ class VariationalArgs: def __init__( self, - algorithm: Optional[str] = None, - iter: Optional[int] = None, - grad_samples: Optional[int] = None, - elbo_samples: Optional[int] = None, - eta: Optional[float] = None, - adapt_iter: Optional[int] = None, + algorithm: str | None = None, + iter: int | None = None, + grad_samples: int | None = None, + elbo_samples: int | None = None, + eta: float | None = None, + adapt_iter: int | None = None, adapt_engaged: bool = True, - tol_rel_obj: Optional[float] = None, - eval_elbo: Optional[int] = None, - output_samples: Optional[int] = None, + tol_rel_obj: float | None = None, + eval_elbo: int | None = None, + output_samples: int | None = None, ) -> None: self.algorithm = algorithm self.iter = iter @@ -567,7 +567,7 @@ def __init__( def validate( self, - chains: Optional[int] = None, # pylint: disable=unused-argument + chains: int | None = None, # pylint: disable=unused-argument ) -> None: """ Check arguments correctness and consistency. @@ -633,23 +633,23 @@ def __init__( self, model_name: str, model_exe: str, - chain_ids: Optional[list[int]], - method_args: Union[ - SamplerArgs, - OptimizeArgs, - GenerateQuantitiesArgs, - VariationalArgs, - LaplaceArgs, - PathfinderArgs, - ], - data: Union[Mapping[str, Any], str, None] = None, - seed: Union[int, np.integer, list[int], list[np.integer], None] = None, - inits: Union[int, float, str, list[str], None] = None, + chain_ids: list[int] | None, + method_args: ( + SamplerArgs + | OptimizeArgs + | GenerateQuantitiesArgs + | VariationalArgs + | LaplaceArgs + | PathfinderArgs + ), + data: Mapping[str, Any] | str | None = None, + seed: int | np.integer | list[int] | list[np.integer] | None = None, + inits: int | float | str | list[str] | None = None, output_dir: OptionalPath = None, - sig_figs: Optional[int] = None, + sig_figs: int | None = None, save_latent_dynamics: bool = False, save_profile: bool = False, - refresh: Optional[int] = None, + refresh: int | None = None, ) -> None: """Initialize object.""" self.model_name = model_name @@ -839,8 +839,8 @@ def compose_command( idx: int, csv_file: str, *, - diagnostic_file: Optional[str] = None, - profile_file: Optional[str] = None, + diagnostic_file: str | None = None, + profile_file: str | None = None, ) -> list[str]: """ Compose CmdStan command for non-default arguments. diff --git a/cmdstanpy/compilation.py b/cmdstanpy/compilation.py index 6a48563c..ae3a857d 100644 --- a/cmdstanpy/compilation.py +++ b/cmdstanpy/compilation.py @@ -10,7 +10,7 @@ import subprocess from datetime import datetime from pathlib import Path -from typing import Any, Iterable, Optional, Union +from typing import Any, Iterable from cmdstanpy.utils import get_logger from cmdstanpy.utils.cmdstan import ( @@ -57,7 +57,7 @@ 'version', ] -OptionalPath = Union[str, os.PathLike, None] +OptionalPath = str | os.PathLike | None class CompilerOptions: @@ -73,8 +73,8 @@ class CompilerOptions: def __init__( self, *, - stanc_options: Optional[dict[str, Any]] = None, - cpp_options: Optional[dict[str, Any]] = None, + stanc_options: dict[str, Any] | None = None, + cpp_options: dict[str, Any] | None = None, user_header: OptionalPath = None, ) -> None: """Initialize object.""" @@ -88,14 +88,14 @@ def __repr__(self) -> str: ) @property - def stanc_options(self) -> dict[str, Union[bool, int, str, Iterable[str]]]: + def stanc_options(self) -> dict[str, bool | int | str | Iterable[str]]: """Stanc compiler options.""" return self._stanc_options @property - def cpp_options(self) -> dict[str, Union[bool, int]]: + def cpp_options(self) -> dict[str, bool | int]: """C++ compiler options.""" - return self._cpp_options + return self._cpp_options # type: ignore @property def user_header(self) -> str: @@ -220,7 +220,7 @@ def validate_user_header(self) -> None: self._cpp_options['USER_HEADER'] = self._user_header - def compose_stanc(self, filename_in_msg: Optional[str]) -> list[str]: + def compose_stanc(self, filename_in_msg: str | None) -> list[str]: opts = [] if filename_in_msg is not None: @@ -244,7 +244,7 @@ def compose_stanc(self, filename_in_msg: Optional[str]) -> list[str]: opts.append(f'--{key}') return opts - def compose(self, filename_in_msg: Optional[str] = None) -> list[str]: + def compose(self, filename_in_msg: str | None = None) -> list[str]: """ Format makefile options as list of strings. @@ -266,7 +266,7 @@ def compose(self, filename_in_msg: Optional[str] = None) -> list[str]: def src_info( stan_file: str, - stanc_options: Optional[dict[str, Any]] = None, + stanc_options: dict[str, Any] | None = None, ) -> dict[str, Any]: """ Get source info for Stan program file. @@ -291,10 +291,10 @@ def src_info( def compile_stan_file( - src: Union[str, Path], + src: str | Path, force: bool = False, - stanc_options: Optional[dict[str, Any]] = None, - cpp_options: Optional[dict[str, Any]] = None, + stanc_options: dict[str, Any] | None = None, + cpp_options: dict[str, Any] | None = None, user_header: OptionalPath = None, ) -> str: """ @@ -423,13 +423,13 @@ def compile_stan_file( def format_stan_file( - stan_file: Union[str, os.PathLike], + stan_file: str | os.PathLike, *, overwrite_file: bool = False, - canonicalize: Union[bool, str, Iterable[str]] = False, + canonicalize: bool | str | Iterable[str] = False, max_line_length: int = 78, backup: bool = True, - stanc_options: Optional[dict[str, Any]] = None, + stanc_options: dict[str, Any] | None = None, ) -> None: """ Run stanc's auto-formatter on the model code. Either saves directly diff --git a/cmdstanpy/install_cmdstan.py b/cmdstanpy/install_cmdstan.py index 6eb2fde7..06be41f9 100644 --- a/cmdstanpy/install_cmdstan.py +++ b/cmdstanpy/install_cmdstan.py @@ -32,7 +32,7 @@ from functools import cached_property from pathlib import Path from time import sleep -from typing import Any, Callable, Optional, Union +from typing import Any, Callable from tqdm.auto import tqdm @@ -123,8 +123,8 @@ class InstallationSettings: def __init__( self, *, - version: Optional[str] = None, - dir: Optional[str] = None, + version: str | None = None, + dir: str | None = None, progress: bool = False, verbose: bool = False, overwrite: bool = False, @@ -301,7 +301,7 @@ def build(verbose: bool = False, progress: bool = True, cores: int = 1) -> None: @progbar.wrap_callback -def _wrap_build_progress_hook() -> Optional[Callable[[str], None]]: +def _wrap_build_progress_hook() -> Callable[[str], None] | None: """Sets up tqdm callback for CmdStan sampler console msgs.""" pad = ' ' * 20 msgs_expected = 150 # hack: 2.27 make build send ~140 msgs to console @@ -475,9 +475,9 @@ def retrieve_version(version: str, progress: bool = True) -> None: for i in range(6): # always retry to allow for transient URLErrors try: if progress and progbar.allow_show_progress(): - progress_hook: Optional[ - Callable[[int, int, int], None] - ] = wrap_url_progress_hook() + progress_hook: Callable[[int, int, int], None] | None = ( + wrap_url_progress_hook() + ) else: progress_hook = None file_tmp, _ = urllib.request.urlretrieve( @@ -579,7 +579,7 @@ def run_compiler_install(dir: str, verbose: bool, progress: bool) -> None: cxx_toolchain_path(cxx_version, dir) -def run_install(args: Union[InteractiveSettings, InstallationSettings]) -> None: +def run_install(args: InteractiveSettings | InstallationSettings) -> None: """ Run a (potentially interactive) installation """ diff --git a/cmdstanpy/model.py b/cmdstanpy/model.py index f85321ef..9ae7f9ef 100644 --- a/cmdstanpy/model.py +++ b/cmdstanpy/model.py @@ -13,7 +13,7 @@ from concurrent.futures import ThreadPoolExecutor from io import StringIO from multiprocessing import cpu_count -from typing import Any, Callable, Mapping, Optional, Sequence, TypeVar, Union +from typing import Any, Callable, Mapping, Sequence, TypeVar import numpy as np import pandas as pd @@ -55,7 +55,7 @@ from . import progress as progbar -OptionalPath = Union[str, os.PathLike, None] +OptionalPath = str | os.PathLike | None Fit = TypeVar('Fit', CmdStanMCMC, CmdStanMLE, CmdStanVB) @@ -96,8 +96,8 @@ def __init__( stan_file: OptionalPath = None, exe_file: OptionalPath = None, force_compile: bool = False, - stanc_options: Optional[dict[str, Any]] = None, - cpp_options: Optional[dict[str, Any]] = None, + stanc_options: dict[str, Any] | None = None, + cpp_options: dict[str, Any] | None = None, user_header: OptionalPath = None, ) -> None: """ @@ -242,7 +242,7 @@ def src_info(self) -> dict[str, Any]: return {} return compilation.src_info(str(self.stan_file), self._stanc_options) - def code(self) -> Optional[str]: + def code(self) -> str | None: """Return Stan program as a string.""" if not self._stan_file: raise RuntimeError('Please specify source file') @@ -259,27 +259,27 @@ def code(self) -> Optional[str]: def optimize( self, - data: Union[Mapping[str, Any], str, os.PathLike, None] = None, - seed: Optional[int] = None, - inits: Union[Mapping[str, Any], float, str, os.PathLike, None] = None, + data: Mapping[str, Any] | str | os.PathLike | None = None, + seed: int | None = None, + inits: Mapping[str, Any] | float | str | os.PathLike | None = None, output_dir: OptionalPath = None, - sig_figs: Optional[int] = None, + sig_figs: int | None = None, save_profile: bool = False, - algorithm: Optional[str] = None, - init_alpha: Optional[float] = None, - tol_obj: Optional[float] = None, - tol_rel_obj: Optional[float] = None, - tol_grad: Optional[float] = None, - tol_rel_grad: Optional[float] = None, - tol_param: Optional[float] = None, - history_size: Optional[int] = None, - iter: Optional[int] = None, + algorithm: str | None = None, + init_alpha: float | None = None, + tol_obj: float | None = None, + tol_rel_obj: float | None = None, + tol_grad: float | None = None, + tol_rel_grad: float | None = None, + tol_param: float | None = None, + history_size: int | None = None, + iter: int | None = None, save_iterations: bool = False, require_converged: bool = True, show_console: bool = False, - refresh: Optional[int] = None, + refresh: int | None = None, time_fmt: str = "%Y%m%d%H%M%S", - timeout: Optional[float] = None, + timeout: float | None = None, jacobian: bool = False, # would be nice to move this further up, but that's a breaking change ) -> CmdStanMLE: @@ -451,50 +451,50 @@ def optimize( # pylint: disable=too-many-arguments def sample( self, - data: Union[Mapping[str, Any], str, os.PathLike, None] = None, - chains: Optional[int] = None, - parallel_chains: Optional[int] = None, - threads_per_chain: Optional[int] = None, - seed: Union[int, list[int], None] = None, - chain_ids: Union[int, list[int], None] = None, - inits: Union[ - Mapping[str, Any], - float, - str, - Sequence[Union[str, Mapping[str, Any]]], - None, - ] = None, - iter_warmup: Optional[int] = None, - iter_sampling: Optional[int] = None, + data: Mapping[str, Any] | str | os.PathLike | None = None, + chains: int | None = None, + parallel_chains: int | None = None, + threads_per_chain: int | None = None, + seed: int | list[int] | None = None, + chain_ids: int | list[int] | None = None, + inits: ( + Mapping[str, Any] + | float + | str + | Sequence[str | Mapping[str, Any]] + | None + ) = None, + iter_warmup: int | None = None, + iter_sampling: int | None = None, save_warmup: bool = False, - thin: Optional[int] = None, - max_treedepth: Optional[int] = None, - metric: Optional[str] = None, - step_size: Union[float, list[float], None] = None, + thin: int | None = None, + max_treedepth: int | None = None, + metric: str | None = None, + step_size: float | list[float] | None = None, adapt_engaged: bool = True, - adapt_delta: Optional[float] = None, - adapt_init_phase: Optional[int] = None, - adapt_metric_window: Optional[int] = None, - adapt_step_size: Optional[int] = None, + adapt_delta: float | None = None, + adapt_init_phase: int | None = None, + adapt_metric_window: int | None = None, + adapt_step_size: int | None = None, fixed_param: bool = False, - output_dir: OptionalPath = None, - sig_figs: Optional[int] = None, + output_dir: OptionalPath = None, # Update OptionalPath if needed + sig_figs: int | None = None, save_latent_dynamics: bool = False, save_profile: bool = False, show_progress: bool = True, show_console: bool = False, - refresh: Optional[int] = None, + refresh: int | None = None, time_fmt: str = "%Y%m%d%H%M%S", - timeout: Optional[float] = None, + timeout: float | None = None, *, - force_one_process_per_chain: Optional[bool] = None, - inv_metric: Union[ - str, - np.ndarray, - Mapping[str, Any], - Sequence[Union[str, np.ndarray, Mapping[str, Any]]], - None, - ] = None, + force_one_process_per_chain: bool | None = None, + inv_metric: ( + str + | np.ndarray + | Mapping[str, Any] + | Sequence[str | np.ndarray | Mapping[str, Any]] + | None + ) = None, ) -> CmdStanMCMC: """ Run or more chains of the NUTS-HMC sampler to produce a set of draws @@ -813,8 +813,8 @@ def sample( temp_inits(inits, id=chain_ids[0]) as _inits, temp_metrics(inv_metric, id=chain_ids[0]) as _inv_metric, ): - cmdstan_inits: Union[str, list[str], int, float, None] - cmdstan_metrics: Union[str, list[str], None] + cmdstan_inits: str | list[str] | int | float | None + cmdstan_metrics: str | list[str] | None if one_process_per_chain and isinstance(inits, list): # legacy cmdstan_inits = [ @@ -869,7 +869,7 @@ def sample( show_progress = show_progress and progbar.allow_show_progress() get_logger().info('CmdStan start processing') - progress_hook: Optional[Callable[[str, int], None]] = None + progress_hook: Callable[[str, int], None] | None = None if show_progress: iter_total = 0 if iter_warmup is None: @@ -957,15 +957,15 @@ def sample( def generate_quantities( self, - data: Union[Mapping[str, Any], str, os.PathLike, None] = None, - previous_fit: Union[Fit, list[str], None] = None, - seed: Optional[int] = None, + data: Mapping[str, Any] | str | os.PathLike | None = None, + previous_fit: Fit | list[str] | None = None, + seed: int | None = None, gq_output_dir: OptionalPath = None, - sig_figs: Optional[int] = None, + sig_figs: int | None = None, show_console: bool = False, - refresh: Optional[int] = None, + refresh: int | None = None, time_fmt: str = "%Y%m%d%H%M%S", - timeout: Optional[float] = None, + timeout: float | None = None, ) -> CmdStanGQ[Fit]: """ Run CmdStan's generate_quantities method which runs the generated @@ -1042,7 +1042,7 @@ def generate_quantities( ) try: fit_csv_files = previous_fit - fit_object = from_csv(fit_csv_files) # type: ignore + fit_object: Fit = from_csv(fit_csv_files) # type: ignore except ValueError as e: raise ValueError( 'Invalid sample from Stan CSV files, error:\n\t{}\n\t' @@ -1135,28 +1135,28 @@ def generate_quantities( def variational( self, - data: Union[Mapping[str, Any], str, os.PathLike, None] = None, - seed: Optional[int] = None, - inits: Optional[float] = None, + data: Mapping[str, Any] | str | os.PathLike | None = None, + seed: int | None = None, + inits: float | None = None, output_dir: OptionalPath = None, - sig_figs: Optional[int] = None, + sig_figs: int | None = None, save_latent_dynamics: bool = False, save_profile: bool = False, - algorithm: Optional[str] = None, - iter: Optional[int] = None, - grad_samples: Optional[int] = None, - elbo_samples: Optional[int] = None, - eta: Optional[float] = None, + algorithm: str | None = None, + iter: int | None = None, + grad_samples: int | None = None, + elbo_samples: int | None = None, + eta: float | None = None, adapt_engaged: bool = True, - adapt_iter: Optional[int] = None, - tol_rel_obj: Optional[float] = None, - eval_elbo: Optional[int] = None, - draws: Optional[int] = None, + adapt_iter: int | None = None, + tol_rel_obj: float | None = None, + eval_elbo: int | None = None, + draws: int | None = None, require_converged: bool = True, show_console: bool = False, - refresh: Optional[int] = None, + refresh: int | None = None, time_fmt: str = "%Y%m%d%H%M%S", - timeout: Optional[float] = None, + timeout: float | None = None, ) -> CmdStanVB: """ Run CmdStan's variational inference algorithm to approximate @@ -1341,39 +1341,39 @@ def variational( def pathfinder( self, - data: Union[Mapping[str, Any], str, os.PathLike, None] = None, + data: Mapping[str, Any] | str | os.PathLike | None = None, *, - init_alpha: Optional[float] = None, - tol_obj: Optional[float] = None, - tol_rel_obj: Optional[float] = None, - tol_grad: Optional[float] = None, - tol_rel_grad: Optional[float] = None, - tol_param: Optional[float] = None, - history_size: Optional[int] = None, - num_paths: Optional[int] = None, - max_lbfgs_iters: Optional[int] = None, - draws: Optional[int] = None, - num_single_draws: Optional[int] = None, - num_elbo_draws: Optional[int] = None, + init_alpha: float | None = None, + tol_obj: float | None = None, + tol_rel_obj: float | None = None, + tol_grad: float | None = None, + tol_rel_grad: float | None = None, + tol_param: float | None = None, + history_size: int | None = None, + num_paths: int | None = None, + max_lbfgs_iters: int | None = None, + draws: int | None = None, + num_single_draws: int | None = None, + num_elbo_draws: int | None = None, psis_resample: bool = True, calculate_lp: bool = True, # arguments standard to all methods - seed: Optional[int] = None, - inits: Union[ - Mapping[str, Any], - float, - str, - Sequence[Union[str, Mapping[str, Any]]], - None, - ] = None, + seed: int | None = None, + inits: ( + Mapping[str, Any] + | float + | str + | Sequence[str | Mapping[str, Any]] + | None + ) = None, output_dir: OptionalPath = None, - sig_figs: Optional[int] = None, + sig_figs: int | None = None, save_profile: bool = False, show_console: bool = False, - refresh: Optional[int] = None, + refresh: int | None = None, time_fmt: str = "%Y%m%d%H%M%S", - timeout: Optional[float] = None, - num_threads: Optional[int] = None, + timeout: float | None = None, + num_threads: int | None = None, ) -> CmdStanPathfinder: """ Run CmdStan's Pathfinder variational inference algorithm. @@ -1576,11 +1576,11 @@ def pathfinder( def log_prob( self, - params: Union[dict[str, Any], str, os.PathLike], - data: Union[Mapping[str, Any], str, os.PathLike, None] = None, + params: dict[str, Any] | str | os.PathLike, + data: Mapping[str, Any] | str | os.PathLike | None = None, *, jacobian: bool = True, - sig_figs: Optional[int] = None, + sig_figs: int | None = None, ) -> pd.DataFrame: """ Calculate the log probability and gradient at the given parameter @@ -1659,20 +1659,20 @@ def log_prob( def laplace_sample( self, - data: Union[Mapping[str, Any], str, os.PathLike, None] = None, - mode: Union[CmdStanMLE, str, os.PathLike, None] = None, - draws: Optional[int] = None, + data: Mapping[str, Any] | str | os.PathLike | None = None, + mode: CmdStanMLE | str | os.PathLike | None = None, + draws: int | None = None, *, jacobian: bool = True, # NB: Different than optimize! - seed: Optional[int] = None, + seed: int | None = None, output_dir: OptionalPath = None, - sig_figs: Optional[int] = None, + sig_figs: int | None = None, save_profile: bool = False, show_console: bool = False, - refresh: Optional[int] = None, + refresh: int | None = None, time_fmt: str = "%Y%m%d%H%M%S", - timeout: Optional[float] = None, - opt_args: Optional[dict[str, Any]] = None, + timeout: float | None = None, + opt_args: dict[str, Any] | None = None, ) -> CmdStanLaplace: """ Run a Laplace approximation around the posterior mode. @@ -1816,8 +1816,8 @@ def _run_cmdstan( idx: int, show_progress: bool = False, show_console: bool = False, - progress_hook: Optional[Callable[[str, int], None]] = None, - timeout: Optional[float] = None, + progress_hook: Callable[[str, int], None] | None = None, + timeout: float | None = None, ) -> None: """ Helper function which encapsulates call to CmdStan. @@ -1854,7 +1854,7 @@ def _run_cmdstan( env=os.environ, universal_newlines=True, ) - timer: Optional[threading.Timer] + timer: threading.Timer | None if timeout: def _timer_target() -> None: @@ -1916,7 +1916,7 @@ def _timer_target() -> None: def _wrap_sampler_progress_hook( chain_ids: list[int], total: int, - ) -> Optional[Callable[[str, int], None]]: + ) -> Callable[[str, int], None] | None: """ Sets up tqdm callback for CmdStan sampler console msgs. CmdStan progress messages start with "Iteration", for single chain @@ -1955,13 +1955,13 @@ def progress_hook(line: str, idx: int) -> None: def diagnose( self, - inits: Union[dict[str, Any], str, os.PathLike, None] = None, - data: Union[Mapping[str, Any], str, os.PathLike, None] = None, + inits: dict[str, Any] | str | os.PathLike | None = None, + data: Mapping[str, Any] | str | os.PathLike | None = None, *, - epsilon: Optional[float] = None, - error: Optional[float] = None, + epsilon: float | None = None, + error: float | None = None, require_gradients_ok: bool = True, - sig_figs: Optional[int] = None, + sig_figs: int | None = None, ) -> pd.DataFrame: """ Run diagnostics to calculate the gradients at the specified parameter diff --git a/cmdstanpy/stanfit/__init__.py b/cmdstanpy/stanfit/__init__.py index 35a76ca3..01d84330 100644 --- a/cmdstanpy/stanfit/__init__.py +++ b/cmdstanpy/stanfit/__init__.py @@ -2,7 +2,6 @@ import glob import os -from typing import Optional, Union from cmdstanpy.cmdstan_args import ( CmdStanArgs, @@ -36,11 +35,16 @@ def from_csv( - path: Union[str, list[str], os.PathLike, None] = None, - method: Optional[str] = None, -) -> Union[ - CmdStanMCMC, CmdStanMLE, CmdStanVB, CmdStanPathfinder, CmdStanLaplace, None -]: + path: str | list[str] | os.PathLike | None = None, + method: str | None = None, +) -> ( + CmdStanMCMC + | CmdStanMLE + | CmdStanVB + | CmdStanPathfinder + | CmdStanLaplace + | None +): """ Instantiate a CmdStan object from a the Stan CSV files from a CmdStan run. CSV files are specified from either a list of Stan CSV files or a single diff --git a/cmdstanpy/stanfit/gq.py b/cmdstanpy/stanfit/gq.py index 40f17893..6a1bfc63 100644 --- a/cmdstanpy/stanfit/gq.py +++ b/cmdstanpy/stanfit/gq.py @@ -3,6 +3,8 @@ generate quantities (GQ) method """ +from __future__ import annotations + from collections import Counter from typing import ( Any, @@ -10,9 +12,7 @@ Hashable, MutableMapping, NoReturn, - Optional, TypeVar, - Union, overload, ) @@ -257,7 +257,7 @@ def draws( def draws_pd( self, - vars: Union[list[str], str, None] = None, + vars: list[str] | str | None = None, inc_warmup: bool = False, inc_sample: bool = False, ) -> pd.DataFrame: @@ -401,25 +401,23 @@ def draws_pd( @overload def draws_xr( - self: Union["CmdStanGQ[CmdStanMLE]", "CmdStanGQ[CmdStanVB]"], - vars: Union[str, list[str], None] = None, + self: CmdStanGQ[CmdStanMLE] | CmdStanGQ[CmdStanVB], + vars: str | list[str] | None = None, inc_warmup: bool = False, inc_sample: bool = False, - ) -> NoReturn: - ... + ) -> NoReturn: ... @overload def draws_xr( self: "CmdStanGQ[CmdStanMCMC]", - vars: Union[str, list[str], None] = None, + vars: str | list[str] | None = None, inc_warmup: bool = False, inc_sample: bool = False, - ) -> "xr.Dataset": - ... + ) -> "xr.Dataset": ... def draws_xr( self, - vars: Union[str, list[str], None] = None, + vars: str | list[str] | None = None, inc_warmup: bool = False, inc_sample: bool = False, ) -> "xr.Dataset": @@ -689,7 +687,7 @@ def _previous_draws_pd( self, vars: list[str], inc_warmup: bool ) -> pd.DataFrame: if vars: - sel: Union[list[str], slice] = vars + sel: list[str] | slice = vars else: sel = slice(None, None) @@ -705,7 +703,7 @@ def _previous_draws_pd( else: # CmdStanVB: return p_fit.variational_sample_pd[sel] - def save_csvfiles(self, dir: Optional[str] = None) -> None: + def save_csvfiles(self, dir: str | None = None) -> None: """ Move output CSV files to specified directory. If files were written to the temporary session directory, clean filename. @@ -723,7 +721,7 @@ def save_csvfiles(self, dir: Optional[str] = None) -> None: # TODO(2.0): remove @property - def mcmc_sample(self) -> Union[CmdStanMCMC, CmdStanMLE, CmdStanVB]: + def mcmc_sample(self) -> CmdStanMCMC | CmdStanMLE | CmdStanVB: get_logger().warning( "Property `mcmc_sample` is deprecated, use `previous_fit` instead" ) diff --git a/cmdstanpy/stanfit/laplace.py b/cmdstanpy/stanfit/laplace.py index 7e36d4e7..314cccf6 100644 --- a/cmdstanpy/stanfit/laplace.py +++ b/cmdstanpy/stanfit/laplace.py @@ -2,7 +2,7 @@ Container for the result of running a laplace approximation. """ -from typing import Any, Hashable, MutableMapping, Optional, Union +from typing import Any, Hashable, MutableMapping import numpy as np import pandas as pd @@ -41,8 +41,8 @@ def __init__(self, runset: RunSet, mode: CmdStanMLE) -> None: self._metadata = InferenceMetadata.from_csv(self._runset.csv_files[0]) def create_inits( - self, seed: Optional[int] = None, chains: int = 4 - ) -> Union[list[dict[str, np.ndarray]], dict[str, np.ndarray]]: + self, seed: int | None = None, chains: int = 4 + ) -> list[dict[str, np.ndarray]] | dict[str, np.ndarray]: """ Create initial values for the parameters of the model by randomly selecting draws from the Laplace approximation. @@ -168,7 +168,7 @@ def draws(self) -> np.ndarray: def draws_pd( self, - vars: Union[list[str], str, None] = None, + vars: list[str] | str | None = None, ) -> pd.DataFrame: if vars is not None: if isinstance(vars, str): @@ -197,7 +197,7 @@ def draws_pd( def draws_xr( self, - vars: Union[str, list[str], None] = None, + vars: str | list[str] | None = None, ) -> "xr.Dataset": """ Returns the sampler draws as a xarray Dataset. @@ -308,7 +308,7 @@ def column_names(self) -> tuple[str, ...]: """ return self._metadata.column_names - def save_csvfiles(self, dir: Optional[str] = None) -> None: + def save_csvfiles(self, dir: str | None = None) -> None: """ Move output CSV files to specified directory. If files were written to the temporary session directory, clean filename. diff --git a/cmdstanpy/stanfit/mcmc.py b/cmdstanpy/stanfit/mcmc.py index 7ff7d142..7400db9f 100644 --- a/cmdstanpy/stanfit/mcmc.py +++ b/cmdstanpy/stanfit/mcmc.py @@ -5,7 +5,7 @@ import math import os from io import StringIO -from typing import Any, Hashable, MutableMapping, Optional, Sequence, Union +from typing import Any, Hashable, MutableMapping, Sequence import numpy as np import pandas as pd @@ -97,8 +97,8 @@ def __init__( self._check_sampler_diagnostics() def create_inits( - self, seed: Optional[int] = None, chains: int = 4 - ) -> Union[list[dict[str, np.ndarray]], dict[str, np.ndarray]]: + self, seed: int | None = None, chains: int = 4 + ) -> list[dict[str, np.ndarray]] | dict[str, np.ndarray]: """ Create initial values for the parameters of the model by randomly selecting draws from the MCMC samples. If the samples @@ -211,7 +211,7 @@ def column_names(self) -> tuple[str, ...]: return self._metadata.column_names @property - def metric_type(self) -> Optional[str]: + def metric_type(self) -> str | None: """ Metric type used for adaptation, either 'diag_e' or 'dense_e', according to CmdStan arg 'metric'. @@ -225,7 +225,7 @@ def metric_type(self) -> Optional[str]: # TODO(2.0): remove @property - def metric(self) -> Optional[np.ndarray]: + def metric(self) -> np.ndarray | None: """Deprecated. Use ``.inv_metric`` instead.""" get_logger().warning( 'The "metric" property is deprecated, use "inv_metric" instead. ' @@ -234,7 +234,7 @@ def metric(self) -> Optional[np.ndarray]: return self.inv_metric @property - def inv_metric(self) -> Optional[np.ndarray]: + def inv_metric(self) -> np.ndarray | None: """ Inverse mass matrix used by sampler for each chain. Returns a ``nchains x nparams`` array when metric_type is 'diag_e', @@ -248,7 +248,7 @@ def inv_metric(self) -> Optional[np.ndarray]: return self._metric @property - def step_size(self) -> Optional[np.ndarray]: + def step_size(self) -> np.ndarray | None: """ Step size used by sampler for each chain. When sampler algorithm 'fixed_param' is specified, step size is None. @@ -264,7 +264,7 @@ def thin(self) -> int: return self._thin @property - def divergences(self) -> Optional[np.ndarray]: + def divergences(self) -> np.ndarray | None: """ Per-chain total number of post-warmup divergent iterations. When sampler algorithm 'fixed_param' is specified, returns None. @@ -272,7 +272,7 @@ def divergences(self) -> Optional[np.ndarray]: return self._divergences if not self._is_fixed_param else None @property - def max_treedepths(self) -> Optional[np.ndarray]: + def max_treedepths(self) -> np.ndarray | None: """ Per-chain total number of post-warmup iterations where the NUTS sampler reached the maximum allowed treedepth. @@ -564,7 +564,7 @@ def summary( summary_data.index.name = None return summary_data[mask] - def diagnose(self) -> Optional[str]: + def diagnose(self) -> str | None: """ Run cmdstan/bin/diagnose over all output CSV files, return console output. @@ -586,7 +586,7 @@ def diagnose(self) -> Optional[str]: def draws_pd( self, - vars: Union[list[str], str, None] = None, + vars: list[str] | str | None = None, inc_warmup: bool = False, ) -> pd.DataFrame: """ @@ -664,7 +664,7 @@ def draws_pd( )[cols] def draws_xr( - self, vars: Union[str, list[str], None] = None, inc_warmup: bool = False + self, vars: str | list[str] | None = None, inc_warmup: bool = False ) -> "xr.Dataset": """ Returns the sampler draws as a xarray Dataset. @@ -822,7 +822,7 @@ def method_variables(self) -> dict[str, np.ndarray]: for name, var in self._metadata.method_vars.items() } - def save_csvfiles(self, dir: Optional[str] = None) -> None: + def save_csvfiles(self, dir: str | None = None) -> None: """ Move output CSV files to specified directory. If files were written to the temporary session directory, clean filename. diff --git a/cmdstanpy/stanfit/metadata.py b/cmdstanpy/stanfit/metadata.py index 224b7c59..da322f1e 100644 --- a/cmdstanpy/stanfit/metadata.py +++ b/cmdstanpy/stanfit/metadata.py @@ -2,7 +2,7 @@ import copy import os -from typing import Any, Iterator, Union +from typing import Any, Iterator import stanio @@ -17,7 +17,7 @@ class InferenceMetadata: """ def __init__( - self, config: dict[str, Union[str, int, float, tuple[str, ...]]] + self, config: dict[str, str | int | float | tuple[str, ...]] ) -> None: """Initialize object from CSV headers""" self._cmdstan_config = config @@ -33,7 +33,7 @@ def __init__( @classmethod def from_csv( - cls, stan_csv: Union[str, os.PathLike, Iterator[bytes]] + cls, stan_csv: str | os.PathLike | Iterator[bytes] ) -> 'InferenceMetadata': try: comments, header, _ = stancsv.parse_comments_header_and_draws( @@ -48,7 +48,7 @@ def from_csv( def __repr__(self) -> str: return 'Metadata:\n{}\n'.format(self._cmdstan_config) - def __getitem__(self, key: str) -> Union[str, int, float, tuple[str, ...]]: + def __getitem__(self, key: str) -> str | int | float | tuple[str, ...]: return self._cmdstan_config[key] @property diff --git a/cmdstanpy/stanfit/mle.py b/cmdstanpy/stanfit/mle.py index 18cbc44d..276a5496 100644 --- a/cmdstanpy/stanfit/mle.py +++ b/cmdstanpy/stanfit/mle.py @@ -1,7 +1,6 @@ """Container for the result of running optimization""" from collections import OrderedDict -from typing import Optional import numpy as np import pandas as pd @@ -58,7 +57,7 @@ def __init__(self, runset: RunSet) -> None: self._all_iters: np.ndarray = all_draws def create_inits( - self, seed: Optional[int] = None, chains: int = 4 + self, seed: int | None = None, chains: int = 4 ) -> dict[str, np.ndarray]: """ Create initial values for the parameters of the model @@ -134,7 +133,7 @@ def optimized_params_np(self) -> np.ndarray: return self._mle @property - def optimized_iterations_np(self) -> Optional[np.ndarray]: + def optimized_iterations_np(self) -> np.ndarray | None: """ Returns all saved iterations from the optimizer and final estimate as a numpy.ndarray which contains all optimizer outputs, i.e., @@ -167,7 +166,7 @@ def optimized_params_pd(self) -> pd.DataFrame: return pd.DataFrame([self._mle], columns=self.column_names) @property - def optimized_iterations_pd(self) -> Optional[pd.DataFrame]: + def optimized_iterations_pd(self) -> pd.DataFrame | None: """ Returns all saved iterations from the optimizer and final estimate as a pandas.DataFrame which contains all optimizer outputs, i.e., @@ -294,7 +293,7 @@ def stan_variables( ) return result - def save_csvfiles(self, dir: Optional[str] = None) -> None: + def save_csvfiles(self, dir: str | None = None) -> None: """ Move output CSV files to specified directory. If files were written to the temporary session directory, clean filename. diff --git a/cmdstanpy/stanfit/pathfinder.py b/cmdstanpy/stanfit/pathfinder.py index 0fbbe5ec..8549c78c 100644 --- a/cmdstanpy/stanfit/pathfinder.py +++ b/cmdstanpy/stanfit/pathfinder.py @@ -2,8 +2,6 @@ Container for the result of running Pathfinder. """ -from typing import Optional, Union - import numpy as np from cmdstanpy.cmdstan_args import Method @@ -30,8 +28,8 @@ def __init__(self, runset: RunSet): self._metadata = InferenceMetadata.from_csv(self._runset.csv_files[0]) def create_inits( - self, seed: Optional[int] = None, chains: int = 4 - ) -> Union[list[dict[str, np.ndarray]], dict[str, np.ndarray]]: + self, seed: int | None = None, chains: int = 4 + ) -> list[dict[str, np.ndarray]] | dict[str, np.ndarray]: """ Create initial values for the parameters of the model by randomly selecting draws from the Pathfinder approximation. @@ -216,7 +214,7 @@ def is_resampled(self) -> bool: in (1, 'true') ) - def save_csvfiles(self, dir: Optional[str] = None) -> None: + def save_csvfiles(self, dir: str | None = None) -> None: """ Move output CSV files to specified directory. If files were written to the temporary session directory, clean filename. diff --git a/cmdstanpy/stanfit/runset.py b/cmdstanpy/stanfit/runset.py index a96c0ea8..cbcb7e7c 100644 --- a/cmdstanpy/stanfit/runset.py +++ b/cmdstanpy/stanfit/runset.py @@ -9,7 +9,6 @@ import tempfile from datetime import datetime from time import time -from typing import Optional from cmdstanpy import _TMPDIR from cmdstanpy.cmdstan_args import CmdStanArgs, Method @@ -31,7 +30,7 @@ def __init__( args: CmdStanArgs, chains: int = 1, *, - chain_ids: Optional[list[int]] = None, + chain_ids: list[int] | None = None, time_fmt: str = "%Y%m%d%H%M%S", one_process_per_chain: bool = True, ) -> None: @@ -162,23 +161,29 @@ def cmd(self, idx: int) -> list[str]: return self._args.compose_command( idx, csv_file=self.csv_files[idx], - diagnostic_file=self.diagnostic_files[idx] - if self._args.save_latent_dynamics - else None, - profile_file=self.profile_files[idx] - if self._args.save_profile - else None, + diagnostic_file=( + self.diagnostic_files[idx] + if self._args.save_latent_dynamics + else None + ), + profile_file=( + self.profile_files[idx] if self._args.save_profile else None + ), ) else: return self._args.compose_command( idx, csv_file=self.file_path('.csv'), - diagnostic_file=self.file_path(".csv", extra="-diagnostic") - if self._args.save_latent_dynamics - else None, - profile_file=self.file_path(".csv", extra="-profile") - if self._args.save_profile - else None, + diagnostic_file=( + self.file_path(".csv", extra="-diagnostic") + if self._args.save_latent_dynamics + else None + ), + profile_file=( + self.file_path(".csv", extra="-profile") + if self._args.save_profile + else None + ), ) @property @@ -213,7 +218,7 @@ def profile_files(self) -> list[str]: # pylint: disable=invalid-name def file_path( - self, suffix: str, *, extra: str = "", id: Optional[int] = None + self, suffix: str, *, extra: str = "", id: int | None = None ) -> str: if id is not None: suffix = f"_{id}{suffix}" @@ -256,7 +261,7 @@ def get_err_msgs(self) -> str: msgs.append('\n\t'.join(errors)) return '\n'.join(msgs) - def save_csvfiles(self, dir: Optional[str] = None) -> None: + def save_csvfiles(self, dir: str | None = None) -> None: """ Moves CSV files to specified directory. diff --git a/cmdstanpy/stanfit/vb.py b/cmdstanpy/stanfit/vb.py index f48d2634..dd00fc8f 100644 --- a/cmdstanpy/stanfit/vb.py +++ b/cmdstanpy/stanfit/vb.py @@ -1,7 +1,6 @@ """Container for the results of running autodiff variational inference""" from collections import OrderedDict -from typing import Optional, Union import numpy as np import pandas as pd @@ -53,8 +52,8 @@ def __init__(self, runset: RunSet) -> None: self._variational_sample: np.ndarray = draws_np[1:] def create_inits( - self, seed: Optional[int] = None, chains: int = 4 - ) -> Union[list[dict[str, np.ndarray]], dict[str, np.ndarray]]: + self, seed: int | None = None, chains: int = 4 + ) -> list[dict[str, np.ndarray]] | dict[str, np.ndarray]: """ Create initial values for the parameters of the model by randomly selecting draws from the variational approximation @@ -248,7 +247,7 @@ def variational_sample_pd(self) -> pd.DataFrame: """ return pd.DataFrame(self._variational_sample, columns=self.column_names) - def save_csvfiles(self, dir: Optional[str] = None) -> None: + def save_csvfiles(self, dir: str | None = None) -> None: """ Move output CSV files to specified directory. If files were written to the temporary session directory, clean filename. diff --git a/cmdstanpy/utils/cmdstan.py b/cmdstanpy/utils/cmdstan.py index 2c1951f8..0c27cb6d 100644 --- a/cmdstanpy/utils/cmdstan.py +++ b/cmdstanpy/utils/cmdstan.py @@ -7,7 +7,7 @@ import subprocess import sys from collections import OrderedDict -from typing import Callable, Optional, Union +from typing import Callable from tqdm.auto import tqdm @@ -83,7 +83,7 @@ def validate_dir(install_dir: str) -> None: ) from e -def get_latest_cmdstan(cmdstan_dir: str) -> Optional[str]: +def get_latest_cmdstan(cmdstan_dir: str) -> str | None: """ Given a valid directory path, find all installed CmdStan versions and return highest (i.e., latest) version number. @@ -198,7 +198,7 @@ def cmdstan_path() -> str: return os.path.normpath(cmdstan) -def cmdstan_version() -> Optional[tuple[int, ...]]: +def cmdstan_version() -> tuple[int, ...] | None: """ Parses version string out of CmdStan makefile variable CMDSTAN_VERSION, returns Tuple(Major, minor). @@ -242,7 +242,7 @@ def cmdstan_version() -> Optional[tuple[int, ...]]: def cmdstan_version_before( - major: int, minor: int, info: Optional[dict[str, str]] = None + major: int, minor: int, info: dict[str, str] | None = None ) -> bool: """ Check that CmdStan version is less than Major.minor version. @@ -273,7 +273,7 @@ def cmdstan_version_before( def cxx_toolchain_path( - version: Optional[str] = None, install_dir: Optional[str] = None + version: str | None = None, install_dir: str | None = None ) -> tuple[str, ...]: """ Validate, then activate C++ toolchain directory path. @@ -434,8 +434,8 @@ def cxx_toolchain_path( def install_cmdstan( - version: Optional[str] = None, - dir: Optional[str] = None, + version: str | None = None, + dir: str | None = None, overwrite: bool = False, compiler: bool = False, progress: bool = False, @@ -493,7 +493,7 @@ def install_cmdstan( run_install, ) - args: Union[InstallationSettings, InteractiveSettings] + args: InstallationSettings | InteractiveSettings if interactive: if any( @@ -538,7 +538,7 @@ def install_cmdstan( @progbar.wrap_callback -def wrap_url_progress_hook() -> Optional[Callable[[int, int, int], None]]: +def wrap_url_progress_hook() -> Callable[[int, int, int], None] | None: """Sets up tqdm callback for url downloads.""" pbar: tqdm = tqdm( unit='B', diff --git a/cmdstanpy/utils/command.py b/cmdstanpy/utils/command.py index 6f3568a6..e669c5dd 100644 --- a/cmdstanpy/utils/command.py +++ b/cmdstanpy/utils/command.py @@ -5,7 +5,7 @@ import os import subprocess import sys -from typing import Callable, Optional, TextIO +from typing import Callable, TextIO from .filesystem import pushd from .logging import get_logger @@ -13,10 +13,10 @@ def do_command( cmd: list[str], - cwd: Optional[str] = None, + cwd: str | None = None, *, - fd_out: Optional[TextIO] = sys.stdout, - pbar: Optional[Callable[[str], None]] = None, + fd_out: TextIO | None = sys.stdout, + pbar: Callable[[str], None] | None = None, ) -> None: """ Run command as subprocess, polls process output pipes and diff --git a/cmdstanpy/utils/filesystem.py b/cmdstanpy/utils/filesystem.py index 0cc5ac80..6f62d53e 100644 --- a/cmdstanpy/utils/filesystem.py +++ b/cmdstanpy/utils/filesystem.py @@ -8,7 +8,7 @@ import re import shutil import tempfile -from typing import Any, Iterator, Mapping, Optional, Sequence, Union +from typing import Any, Iterator, Mapping, Sequence import numpy as np @@ -107,8 +107,8 @@ def pushd(new_dir: str) -> Iterator[None]: def _temp_single_json( - data: Union[str, os.PathLike, Mapping[str, Any], None], -) -> Iterator[Optional[str]]: + data: str | os.PathLike | Mapping[str, Any] | None, +) -> Iterator[str | None]: """Context manager for json files.""" if data is None: yield None @@ -131,9 +131,9 @@ def _temp_single_json( def _temp_multiinput( - input: Union[str, os.PathLike, Mapping[str, Any], Sequence[Any], None], + input: str | os.PathLike | Mapping[str, Any] | Sequence[Any] | None, base: int = 1, -) -> Iterator[Optional[str]]: +) -> Iterator[str | None]: if isinstance(input, Sequence) and not isinstance( input, (str, os.PathLike) ): @@ -171,12 +171,17 @@ def _temp_multiinput( @contextlib.contextmanager def temp_metrics( - metrics: Union[ - str, os.PathLike, Mapping[str, Any], np.ndarray, Sequence[Any], None - ], + metrics: ( + str + | os.PathLike + | Mapping[str, Any] + | np.ndarray + | Sequence[Any] + | None + ), *, id: int = 1, -) -> Iterator[Union[str, None]]: +) -> Iterator[str | None]: if isinstance(metrics, dict): if 'inv_metric' not in metrics: raise ValueError('Entry "inv_metric" not found in metric dict.') @@ -201,13 +206,19 @@ def temp_metrics( @contextlib.contextmanager def temp_inits( - inits: Union[ - str, os.PathLike, Mapping[str, Any], float, int, Sequence[Any], None - ], + inits: ( + str + | os.PathLike + | Mapping[str, Any] + | float + | int + | Sequence[Any] + | None + ), *, allow_multiple: bool = True, id: int = 1, -) -> Iterator[Union[str, float, int, None]]: +) -> Iterator[str | float | int | None]: if isinstance(inits, (float, int)): yield inits return diff --git a/cmdstanpy/utils/logging.py b/cmdstanpy/utils/logging.py index f6115901..eaba8953 100644 --- a/cmdstanpy/utils/logging.py +++ b/cmdstanpy/utils/logging.py @@ -6,7 +6,7 @@ import logging import types from contextlib import AbstractContextManager -from typing import Optional, Type +from typing import Type @functools.lru_cache(maxsize=None) @@ -45,9 +45,9 @@ def __enter__(self) -> "ToggleLogging": def __exit__( self, - exc_type: Optional[Type[BaseException]], - exc_value: Optional[BaseException], - traceback: Optional[types.TracebackType], + exc_type: Type[BaseException] | None, + exc_value: BaseException | None, + traceback: types.TracebackType | None, ) -> None: self.logger.disabled = self.prev_state diff --git a/cmdstanpy/utils/stancsv.py b/cmdstanpy/utils/stancsv.py index 0349a4c1..f8aee8bd 100644 --- a/cmdstanpy/utils/stancsv.py +++ b/cmdstanpy/utils/stancsv.py @@ -8,7 +8,7 @@ import os import re import warnings -from typing import Any, Iterator, Mapping, Optional, Sequence, Union +from typing import Any, Iterator, Mapping, Sequence import numpy as np import numpy.typing as npt @@ -17,8 +17,8 @@ def parse_comments_header_and_draws( - stan_csv: Union[str, os.PathLike, Iterator[bytes]], -) -> tuple[list[bytes], Optional[str], list[bytes]]: + stan_csv: str | os.PathLike | Iterator[bytes], +) -> tuple[list[bytes], str | None, list[bytes]]: """Parses lines of a Stan CSV file into comment lines, the header line, and draws lines. @@ -27,7 +27,7 @@ def parse_comments_header_and_draws( def partition_csv( lines: Iterator[bytes], - ) -> tuple[list[bytes], Optional[str], list[bytes]]: + ) -> tuple[list[bytes], str | None, list[bytes]]: comment_lines: list[bytes] = [] draws_lines: list[bytes] = [] header = None @@ -105,7 +105,7 @@ def csv_bytes_list_to_numpy( def parse_hmc_adaptation_lines( comment_lines: list[bytes], -) -> tuple[Optional[float], Optional[npt.NDArray[np.float64]]]: +) -> tuple[float | None, npt.NDArray[np.float64] | None]: """Extracts step size/mass matrix information from the Stan CSV comment lines by parsing the adaptation section. If the diag_e metric is used, the returned mass matrix will be a 1D array of the diagnoal elements, @@ -163,10 +163,10 @@ def extract_key_val_pairs( def parse_config( comment_lines: list[bytes], -) -> dict[str, Union[str, int, float]]: +) -> dict[str, str | int | float]: """Extracts the key=value config settings from Stan CSV comment lines and returns a dictionary.""" - out: dict[str, Union[str, int, float]] = {} + out: dict[str, str | int | float] = {} for key, val in extract_key_val_pairs(comment_lines): if key == 'file': if not val.endswith('csv'): @@ -194,12 +194,12 @@ def parse_header(header: str) -> tuple[str, ...]: def construct_config_header_dict( - comment_lines: list[bytes], header: Optional[str] -) -> dict[str, Union[str, int, float, tuple[str, ...]]]: + comment_lines: list[bytes], header: str | None +) -> dict[str, str | int | float | tuple[str, ...]]: """Extracts config and header info from comment/draws lines parsed from a Stan CSV file.""" config = parse_config(comment_lines) - out: dict[str, Union[str, int, float, tuple[str, ...]]] = {**config} + out: dict[str, str | int | float | tuple[str, ...]] = {**config} if header: out["raw_header"] = header out["column_names"] = parse_header(header) @@ -264,7 +264,7 @@ def is_sneaky_fixed_param(header: str) -> bool: def count_warmup_and_sampling_draws( - stan_csv: Union[str, os.PathLike, Iterator[bytes]], + stan_csv: str | os.PathLike | Iterator[bytes], ) -> tuple[int, int]: """Scans through a Stan CSV file to count the number of lines in the warmup/sampling blocks to determine counts for warmup and sampling draws. @@ -432,7 +432,7 @@ def parse_timing_lines( def check_sampler_csv( - path: Union[str, os.PathLike], + path: str | os.PathLike, iter_sampling: int = _CMDSTAN_SAMPLING, iter_warmup: int = _CMDSTAN_WARMUP, save_warmup: bool = False, @@ -473,8 +473,8 @@ def check_sampler_csv( def parse_sampler_metadata_from_csv( - path: Union[str, os.PathLike], -) -> dict[str, Union[int, float, str, tuple[str, ...], dict[str, float]]]: + path: str | os.PathLike, +) -> dict[str, int | float | str | tuple[str, ...] | dict[str, float]]: """Parses sampling metadata from a given Stan CSV path for a sample run""" try: comments, header, draws = parse_comments_header_and_draws(path) @@ -504,7 +504,7 @@ def parse_sampler_metadata_from_csv( "Sampling": "sampling", "Total": "total", } - addtl: dict[str, Union[int, dict[str, float]]] = { + addtl: dict[str, int | dict[str, float]] = { "draws_warmup": num_warmup, "draws_sampling": num_sampling, "ct_divergences": divs, @@ -569,7 +569,7 @@ def read_rdump_metric(path: str) -> list[int]: return list(metric_dict['inv_metric'].shape) -def rload(fname: str) -> Optional[dict[str, Union[int, float, np.ndarray]]]: +def rload(fname: str) -> dict[str, int | float | np.ndarray] | None: """Parse data and parameter variable values from an R dump format file. This parser only supports the subset of R dump data as described in the "Dump Data Format" section of the CmdStan manual, i.e., @@ -602,7 +602,7 @@ def rload(fname: str) -> Optional[dict[str, Union[int, float, np.ndarray]]]: return data_dict -def parse_rdump_value(rhs: str) -> Union[int, float, np.ndarray]: +def parse_rdump_value(rhs: str) -> int | float | np.ndarray: """Process right hand side of Rdump variable assignment statement. Value is either scalar, vector, or multi-dim structure. Use regex to capture structure values, dimensions. @@ -611,7 +611,7 @@ def parse_rdump_value(rhs: str) -> Union[int, float, np.ndarray]: r'structure\(\s*c\((?P[^)]*)\)' r'(,\s*\.Dim\s*=\s*c\s*\((?P[^)]*)\s*\))?\)' ) - val: Union[int, float, np.ndarray] + val: int | float | np.ndarray try: if rhs.startswith('structure'): parse = pat.match(rhs) @@ -634,13 +634,13 @@ def parse_rdump_value(rhs: str) -> Union[int, float, np.ndarray]: def try_deduce_metric_type( - inv_metric: Union[ - str, - np.ndarray, - Mapping[str, Any], - Sequence[Union[str, np.ndarray, Mapping[str, Any]]], - ], -) -> Optional[str]: + inv_metric: ( + str + | np.ndarray + | Mapping[str, Any] + | Sequence[str | np.ndarray | Mapping[str, Any]] + ), +) -> str | None: """Given a user-supplied metric, try to infer the correct metric type.""" if isinstance(inv_metric, Sequence) and not isinstance( inv_metric, (str, np.ndarray, Mapping) diff --git a/pyproject.toml b/pyproject.toml index 612b3964..dabe44be 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -8,7 +8,7 @@ description = "Python interface to CmdStan" readme = "README.md" license = { text = "BSD-3-Clause" } authors = [{ name = "Stan Dev Team" }] -requires-python = ">=3.9" +requires-python = ">=3.10" dependencies = ["pandas", "numpy>=1.21", "tqdm", "stanio>=0.4.0,<2.0.0"] dynamic = ["version"] classifiers = [ @@ -67,6 +67,7 @@ docs = [ [tool.black] line-length = 80 skip-string-normalization = true +target-version = ['py310', 'py311', 'py312', 'py313', 'py314'] [tool.isort] multi_line_output = 3 diff --git a/test/test_log_prob.py b/test/test_log_prob.py index 046dd3af..5be47d0b 100644 --- a/test/test_log_prob.py +++ b/test/test_log_prob.py @@ -4,7 +4,6 @@ import os import re from test import check_present -from typing import List, Optional import numpy as np import pytest @@ -34,7 +33,7 @@ ], ) def test_lp_good( - sig_figs: Optional[int], expected: List[str], expected_unadjusted: List[str] + sig_figs: int | None, expected: list[str], expected_unadjusted: list[str] ) -> None: model = CmdStanModel(stan_file=BERN_STAN) params = {"theta": 0.34903938392023830482} diff --git a/test/test_model.py b/test/test_model.py index b46a124c..8334560a 100644 --- a/test/test_model.py +++ b/test/test_model.py @@ -6,7 +6,6 @@ import shutil import tempfile from test import check_present -from typing import List from unittest.mock import patch import numpy as np @@ -212,7 +211,7 @@ def test_compile_with_bad_includes(caplog: pytest.LogCaptureFixture) -> None: ], ) def test_compile_with_includes( - caplog: pytest.LogCaptureFixture, stan_file: str, include_paths: List[str] + caplog: pytest.LogCaptureFixture, stan_file: str, include_paths: list[str] ) -> None: getmtime = os.path.getmtime stan_file = os.path.join(DATAFILES_PATH, stan_file) diff --git a/test/test_stancsv.py b/test/test_stancsv.py index 7c648163..a2f4a0e0 100644 --- a/test/test_stancsv.py +++ b/test/test_stancsv.py @@ -4,7 +4,6 @@ import os from pathlib import Path from test import without_import -from typing import List import numpy as np import pytest @@ -94,7 +93,7 @@ def test_csv_bytes_empty() -> None: def test_parse_comments_header_and_draws() -> None: - lines: List[bytes] = [b"# 1\n", b"a\n", b"3\n", b"# 4\n"] + lines: list[bytes] = [b"# 1\n", b"a\n", b"3\n", b"# 4\n"] ( comment_lines, header,