Merge pull request #492 from stan-dev/feature/436-num-chains

Feature/436 num chains

Author: mitzimorris

cmdstanpy/__init__.py

@@ -2,6 +2,7 @@
 """CmdStanPy Module"""
 
 import atexit
+import logging
 import shutil
 import tempfile
 
@@ -27,9 +28,9 @@
 
 def _cleanup_tmpdir() -> None:
     """Force deletion of _TMPDIR."""
-    print('deleting tmpfiles dir: {}'.format(_TMPDIR))
+    logging.getLogger('cmdstanpy').info('deleting tmpfiles dir: %s', _TMPDIR)
     shutil.rmtree(_TMPDIR, ignore_errors=True)
-    print('done')
+    logging.getLogger('cmdstanpy').info('done')
 
 
 atexit.register(_cleanup_tmpdir)

cmdstanpy/cmdstan_args.py

@@ -4,7 +4,7 @@
 import os
 from enum import Enum, auto
 from time import time
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Dict, List, Mapping, Optional, Union
 
 import numpy as np
 from numpy.random import RandomState
@@ -712,7 +712,7 @@ def __init__(
         method_args: Union[
             SamplerArgs, OptimizeArgs, GenerateQuantitiesArgs, VariationalArgs
         ],
-        data: Union[str, Dict[str, Any], None] = None,
+        data: Union[Mapping[str, Any], str, None] = None,
         seed: Union[int, List[int], None] = None,
         inits: Union[int, float, str, List[str], None] = None,
         output_dir: Optional[str] = None,
@@ -811,6 +811,7 @@ def validate(self) -> None:
                     'Argument "sig_figs" must be an integer between 1 and 18,'
                     ' found {}'.format(self.sig_figs)
                 )
+            # TODO: remove at some future release
             if cmdstan_version_before(2, 25):
                 self.sig_figs = None
                 get_logger().warning(
@@ -897,7 +898,8 @@ def compose_command(
         csv_file: str,
         *,
         diagnostic_file: Optional[str] = None,
-        profile_file: Optional[str] = None
+        profile_file: Optional[str] = None,
+        num_chains: Optional[int] = None
     ) -> List[str]:
         """
         Compose CmdStan command for non-default arguments.
@@ -932,13 +934,15 @@ def compose_command(
                 cmd.append('init={}'.format(self.inits[idx]))
         cmd.append('output')
         cmd.append('file={}'.format(csv_file))
-        if diagnostic_file is not None:
+        if diagnostic_file:
             cmd.append('diagnostic_file={}'.format(diagnostic_file))
-        if profile_file is not None:
+        if profile_file:
             cmd.append('profile_file={}'.format(profile_file))
         if self.refresh is not None:
             cmd.append('refresh={}'.format(self.refresh))
         if self.sig_figs is not None:
             cmd.append('sig_figs={}'.format(self.sig_figs))
         cmd = self.method_args.compose(idx, cmd)
+        if num_chains:
+            cmd.append('num_chains={}'.format(num_chains))
         return cmd