Simplify/clean sampler interface & use

book/notebooks/setRNG.ipynb

@@ -38,27 +38,14 @@
     },
     "tags": []
    },
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "surmise v0.4.1.dev238+g87713c9ed.d20260715\n",
-      "scipy   v1.18.0\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "import numpy as np\n",
     "import itertools as it\n",
     "import scipy as sp\n",
     "import scipy.stats as sps\n",
     "\n",
-    "import surmise\n",
-    "from surmise.emulation import emulator\n",
-    "\n",
-    "print(f\"surmise v{surmise.__version__}\")\n",
-    "print(f\"scipy   v{sp.__version__}\")"
+    "import surmise"
    ]
   },
   {
@@ -75,12 +62,24 @@
      "remove-cell"
     ]
    },
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "surmise v0.4.1.dev215+g936454b96.d20260716\n",
+      "scipy   v1.18.0\n"
+     ]
+    }
+   ],
    "source": [
     "# This is to simulate a fresh session.  Generally a user does not have to clear an RNG before setting one.\n",
     "# Therefore we remove this cell from the book's rendering.\n",
     "from surmise._RandomNumberGenerator import RandomNumberGenerator\n",
-    "RandomNumberGenerator()._clear_RNG()"
+    "RandomNumberGenerator()._clear_RNG()\n",
+    "\n",
+    "print(f\"surmise v{surmise.__version__}\")\n",
+    "print(f\"scipy   v{sp.__version__}\")"
    ]
   },
   {
@@ -130,7 +129,7 @@
     ")\n",
     "\n",
     "try:\n",
-    "    emulator(x=x, theta=theta, f=borehole_model(x, theta), method='PCGP')\n",
+    "    surmise.emulator(x=x, theta=theta, f=borehole_model(x, theta), method='PCGP')\n",
     "except Exception as e:\n",
     "    print(e)"
    ]
@@ -193,7 +192,7 @@
     "theta = sps.uniform.rvs(0, 1, size=(50, 4), random_state=data_rng)\n",
     "\n",
     "# surmise uses the surmise RNG under the hood\n",
-    "emu = emulator(x=x, theta=theta, f=borehole_model(x, theta), method='PCGP')\n",
+    "emu = surmise.emulator(x=x, theta=theta, f=borehole_model(x, theta), method='PCGP')\n",
     "pred = emu.predict(x=x, theta=theta)\n",
     "print(\"prediction mean shape:\", pred.mean().shape)"
    ]
@@ -316,7 +315,7 @@
     "    x = sps.uniform.rvs(0, 1, size=(50, 3), random_state=data_rng)\n",
     "    x[:, 2] = x[:, 2] > 0.5\n",
     "    thetas = sps.uniform.rvs(0, 1, size=(15, 4), random_state=data_rng)\n",
-    "    emu = emulator(x=x, theta=thetas, f=borehole_model(x, thetas), method='PCGP')\n",
+    "    emu = surmise.emulator(x=x, theta=thetas, f=borehole_model(x, thetas), method='PCGP')\n",
     "    return emu.predict(x=x, theta=thetas).mean()\n",
     "\n",
     "surmise_1, surmise_2, data_1, data_2 = np.random.SeedSequence(SEED).spawn(4)\n",

src/surmise/__init__.py

@@ -13,7 +13,12 @@
 __author__ = 'Matthew Plumlee, Özge Sürer, Stefan M. Wild, Moses Y-H. Chan'
 __credits__ = 'Northwestern University, Argonne National Laboratory'
 
+# General-use public interface
 from .set_RNG import set_RNG
+from .emulation import emulator
+from .calibration import calibrator
+
+# Advanced API
 from .create_sampler import create_sampler
 
 f_dir = os.path.dirname(os.path.realpath(__file__))

src/surmise/calibrationmethods/directbayes.py

@@ -110,19 +110,23 @@ def draw_func(n):
         return theta0
 
     # Call the sampler
-    if 'sampler' in sampler_args:
-        sampler_name = sampler_args['sampler']
-        # TODO: The sampler name should likely be its own argument
-        # (non-optional?) to the calibrator rather than hiding it in its own set
-        # of arguments.  Why not make sampler_args a single dictionary that
-        # calling code provides to the calibrator?
-        del sampler_args['sampler']
+    specification = copy.deepcopy(sampler_args)
+    if 'sampler' in specification:
+        sampler_name = specification['sampler']
+        del specification['sampler']
     else:
         sampler_name = 'metropolis_hastings'
-    sampler = create_sampler(sampler_name, sampler_args)
+
+    expert_mode = False
+    if 'expertMode' in specification:
+        expert_mode = specification['expertMode']
+        del specification['expertMode']
+
+    sampler = create_sampler(sampler_name, expert_mode=expert_mode)
     results = sampler(logpost_func=logpostfull,
                       draw_func=draw_func,
-                      scipy_stats_rng=global_RNG)
+                      scipy_stats_rng=global_RNG,
+                      specification=specification)
     theta = results["theta"]
 
     # Update fitinfo dict

src/surmise/calibrationmethods/directbayeswoodbury.py

@@ -151,16 +151,24 @@ def draw_func(n):
 
         return theta0
 
-    # obtain theta draws from posterior distribution
-    if 'sampler' in sampler_args:
-        sampler_name = sampler_args['sampler']
-        del sampler_args['sampler']
+    # Call the sampler
+    specification = copy.deepcopy(sampler_args)
+    if 'sampler' in specification:
+        sampler_name = specification['sampler']
+        del specification['sampler']
     else:
         sampler_name = 'metropolis_hastings'
-    sampler = create_sampler(sampler_name, sampler_args)
+
+    expert_mode = False
+    if 'expertMode' in specification:
+        expert_mode = specification['expertMode']
+        del specification['expertMode']
+
+    sampler = create_sampler(sampler_name, expert_mode=expert_mode)
     results = sampler(logpost_func=logpostfull_wgrad,
                       draw_func=draw_func,
-                      scipy_stats_rng=global_RNG)
+                      scipy_stats_rng=global_RNG,
+                      specification=specification)
     theta = results["theta"]
 
     # obtain log-posterior of theta values

src/surmise/calibrationmethods/mlbayeswoodbury.py

@@ -189,16 +189,24 @@ def draw_func(n):
 
         return theta0
 
-    # obtain theta draws from posterior distribution
-    if 'sampler' in sampler_args:
-        sampler_name = sampler_args['sampler']
-        del sampler_args['sampler']
+    # Call the sampler
+    specification = copy.deepcopy(sampler_args)
+    if 'sampler' in specification:
+        sampler_name = specification['sampler']
+        del specification['sampler']
     else:
         sampler_name = 'metropolis_hastings'
-    sampler = create_sampler(sampler_name, sampler_args)
+
+    expert_mode = False
+    if 'expertMode' in specification:
+        expert_mode = specification['expertMode']
+        del specification['expertMode']
+
+    sampler = create_sampler(sampler_name, expert_mode=expert_mode)
     results = sampler(logpost_func=logpostfull_wgrad,
                       draw_func=draw_func,
-                      scipy_stats_rng=global_RNG)
+                      scipy_stats_rng=global_RNG,
+                      specification=specification)
     theta = results["theta"]
 
     # obtain log-posterior of theta values

src/surmise/calibrationmethods/simulationpost.py

@@ -146,15 +146,24 @@ def draw_func(n):
 
         return theta0
 
-    if 'sampler' in sampler_args:
-        sampler_name = sampler_args['sampler']
-        del sampler_args['sampler']
+    # Call the sampler
+    specification = copy.deepcopy(sampler_args)
+    if 'sampler' in specification:
+        sampler_name = specification['sampler']
+        del specification['sampler']
     else:
         sampler_name = 'metropolis_hastings'
-    sampler = create_sampler(sampler_name, sampler_args)
+
+    expert_mode = False
+    if 'expertMode' in specification:
+        expert_mode = specification['expertMode']
+        del specification['expertMode']
+
+    sampler = create_sampler(sampler_name, expert_mode=expert_mode)
     results = sampler(logpost_func=logpostfull_wgrad,
                       draw_func=draw_func,
-                      scipy_stats_rng=global_RNG)
+                      scipy_stats_rng=global_RNG,
+                      specification=specification)
     theta = results["theta"]
 
     # obtain log-posterior of theta values

src/surmise/create_sampler.py

@@ -1,79 +1,68 @@
-import copy
 import warnings
-import functools
 
 from .utilitiesmethods.metropolis_hastings import sampler as sample_with_metropolis_hastings
 from .utilitiesmethods.LMC import sampler as sample_with_LMC
 from .utilitiesmethods.PTLMC import sampler as sample_with_PTLMC
 
 
-def create_sampler(sampler, options):
+def create_sampler(sampler, expert_mode):
     """
     Construct a sampler function for direct use by |surmise| calibrators.  The
     following example demonstrates its use.
 
     .. code-block:: python
 
-        sample_with_PTLMC = surmise.create_sampler("PTLMC", ptlmc_args)
+        sample_with_PTLMC = surmise.create_sampler("PTLMC", expert_mode=False)
         results = sample_with_PTLMC(
             logpost_func=log_posterior,
             draw_func=draw_from_start_distribution,
-            scipy_stats_rng=np.random.default_rng(RAND_SEED)
+            scipy_stats_rng=np.random.default_rng(RAND_SEED),
+            specification=ptlmc_spec
         )
 
     For typical use cases, samplers are created automatically under-the-hood on
     behalf of users.  Therefore, there is generally no need to explicitly create
     or access samplers.  This function is in the |surmise| public interface only
     as an advanced feature for use by developers and power users.
 
-    .. todo::
-        * The current implementation prevents users from providing a custom
-          sampler to calibrators.  Consider allowing ``sampler`` to be a
-          user-provided sampler function that we assume has the necessary
-          interface.  This function could then confirm that options is ``None``
-          or an empty ``dict`` and just pass that function along.  If the
-          calibrator interface is updated so that calling code must provide a
-          sampler identifier, then that argument could also be setup in this
-          same way.
-
     Parameters
     ----------
     sampler :
         Name of desired sampler offered by |surmise|
-    options :
-        ``dict`` of sampler-specific arguments that fully characterize the
-        desired sampler.  Refer to the documentation of each sampler for more
-        information.
+    expert_mode :
+        Allow the use of research-grade samplers if ``True``
 
     Returns
     -------
     :
         The desired sampler function.
     """
-    KEY = "expertMode"
-
-    if sampler.lower() == "metropolis_hastings":
-        return functools.partial(sample_with_metropolis_hastings, **options)
-    elif sampler.upper() == "LMC":
-        lmc_options = copy.deepcopy(options)
-
-        if KEY in lmc_options:
-            if not isinstance(lmc_options[KEY], bool):
-                raise ValueError(f"{KEY} value must be a boolean")
-            elif not lmc_options[KEY]:
+    if isinstance(sampler, str):
+        if sampler.lower() == "metropolis_hastings":
+            return sample_with_metropolis_hastings
+        elif sampler.upper() == "LMC":
+            if not expert_mode:
                 msg = "{} is included for unofficial research purposes only"
                 raise ValueError(msg.format(sampler))
 
-            del lmc_options[KEY]
-        else:
-            msg = "{} is included for unofficial research purposes only"
-            raise ValueError(msg.format(sampler))
+            # Emit warning to extend a helping hand to the experts.
+            msg = f"Using unofficial research {sampler} sampler"
+            warnings.warn(msg)
+            return sample_with_LMC
+        elif sampler.upper() == "PTLMC":
+            return sample_with_PTLMC
+    elif isinstance(sampler, dict):
+        if len(sampler) != 1:
+            return ValueError('Custom sampler must be {"user": my_sampler_fcn}')
+        source = sampler.keys()[0]
+        if source.lower() != "user":
+            return ValueError('Custom sampler must be {"user": my_sampler_fcn}')
+        sampler_fcn = sampler[source]
+        if not callable(sampler_fcn):
+            return ValueError("Custom sampler function is not callable")
 
-        # Emit warning to extend a helping hand to the experts.
-        msg = f"Using unofficial research {sampler} sampler"
-        warnings.warn(msg)
-        return functools.partial(sample_with_LMC, **lmc_options)
-    elif sampler.upper() == "PTLMC":
-        return functools.partial(sample_with_PTLMC, **options)
+        raise NotImplementedError("This functionality is not under test")
+    else:
+        raise TypeError(f"Sampler should be a string or dict ({sampler})")
 
-    raise TypeError(f"Invalid sampler ({sampler})")
+    raise ValueError(f"Invalid sampler ({sampler})")

src/surmise/utilitiesmethods/LMC.py

@@ -44,8 +44,7 @@
 def sampler(logpost_func,
             draw_func,
             scipy_stats_rng,
-            numsamp=2000,
-            theta0=None):
+            specification):
     '''
 
     Parameters
@@ -73,6 +72,20 @@ def sampler(logpost_func,
         numsamp by p of sampled parameter values
 
     '''
+    VALID_SPECS = {"nSamples", "theta0", "verbose"}
+
+    # Get specification values
+    # TODO: Error check these with useful error messages
+    assert set(specification) == VALID_SPECS
+    numsamp = specification["nSamples"]
+    theta0 = specification["theta0"]
+    verbose = specification["verbose"]
+
+    if verbose:
+        # Don't log theta0 as it could potentially be an overwhelming amount of
+        # information.
+        print(f"nSamples = {numsamp}")
+
     # random number generator
     if not isinstance(scipy_stats_rng, np.random.Generator):
         raise TypeError("Given RNG is not a valid scipy.stats RNG")