import warnings
import importlib
from types import MethodType
from functools import partial, wraps
from multiprocess import Pool, cpu_count
import numpy as np
import SALib.sample as samplers
import SALib.analyze as analyzers
from .util_funcs import ( # noqa: F401, E402
avail_approaches,
_define_problem_with_groups,
) # noqa: F401, E402
from .results import ResultDict
ptqdm_available = True
try:
from p_tqdm import p_imap
except ImportError:
ptqdm_available = False
try:
from pathos.pp import ParallelPythonPool as pp_Pool
pathos_available = True
except ImportError:
pathos_available = False
__all__ = ["ProblemSpec"]
[docs]
class ProblemSpec(dict):
"""Dictionary-like object representing an SALib Problem specification.
Attributes
----------
samples : np.array, of generated samples
results : np.array, of evaluations (i.e., model outputs)
analysis : np.array or dict, of sensitivity indices
"""
def __init__(self, *args, **kwargs):
super(ProblemSpec, self).__init__(*args, **kwargs)
_check_spec_attributes(self)
self._samples = None
self._results = None
self._analysis = None
self["num_vars"] = len(self["names"])
self = _define_problem_with_groups(self)
self._add_samplers()
self._add_analyzers()
@property
def samples(self):
return self._samples
@samples.setter
def samples(self, vals):
cols = vals.shape[1]
if cols != self["num_vars"]:
msg = "Mismatched sample size: Expected "
msg += "{} cols, got {}".format(self["num_vars"], cols)
raise ValueError(msg)
self._samples = vals
# Clear results to avoid confusion
self._results = None
@property
def results(self):
return self._results
@results.setter
def results(self, vals):
val_shape = vals.shape
if len(val_shape) == 1:
cols = 1
else:
cols = vals.shape[1]
out_cols = self.get("outputs", None)
if out_cols is None:
if cols == 1:
self["outputs"] = ["Y"]
else:
self["outputs"] = [f"Y{i}" for i in range(1, cols + 1)]
else:
if cols != len(self["outputs"]):
msg = "Mismatched sample size: Expected "
msg += "{} cols, got {}".format(self["outputs"], cols)
raise ValueError(msg)
self._results = vals
@property
def analysis(self):
return self._analysis
[docs]
def sample(self, func, *args, **kwargs):
"""Create sample using given function.
Parameters
----------
func : function,
Sampling method to use. The given function must accept the SALib
problem specification as the first parameter and return a numpy
array.
*args : list,
Additional arguments to be passed to `func`
**kwargs : dict,
Additional keyword arguments passed to `func`
Returns
-------
self : ProblemSpec object
"""
# Clear model output and analysis results to avoid confusion
# especially if samples are forcibly changed...
self._analysis = None
self._results = None
self._samples = func(self, *args, **kwargs)
return self
[docs]
def set_samples(self, samples: np.ndarray):
"""Set previous samples used."""
self.samples = samples
return self
[docs]
def set_results(self, results: np.ndarray):
"""Set previously available model results."""
if self.samples is not None:
assert self.samples.shape[0] == results.shape[0], (
"Provided result array does not match existing number of"
" existing samples!"
)
self.results = results
return self
[docs]
def evaluate(self, func, *args, **kwargs):
"""Evaluate a given model.
Parameters
----------
func : function,
Model, or function that wraps a model, to be run/evaluated.
The provided function is required to accept a numpy array of
inputs as its first parameter and must return a numpy array of
results.
*args : list,
Additional arguments to be passed to `func`
nprocs : int,
If specified, attempts to parallelize model evaluations
**kwargs : dict,
Additional keyword arguments passed to `func`
Returns
-------
self : ProblemSpec object
"""
nprocs = kwargs.pop("nprocs", 1)
if nprocs > 1:
return self.evaluate_parallel(func, *args, nprocs=nprocs, **kwargs)
self.results = func(self._samples, *args, **kwargs)
return self
[docs]
def evaluate_parallel(self, func, *args, nprocs=None, **kwargs):
"""Evaluate model locally in parallel.
All detected processors will be used if `nprocs` is not specified.
Parameters
----------
func : function,
Model, or function that wraps a model, to be run in parallel.
The provided function needs to accept a numpy array of inputs as
its first parameter and must return a numpy array of results.
nprocs : int,
Number of processors to use.
Capped to the number of available processors.
*args : list,
Additional arguments to be passed to `func`
**kwargs : dict,
Additional keyword arguments passed to `func`
Returns
-------
self : ProblemSpec object
"""
warnings.warn(
"Parallel evaluation is an experimental feature and may not work."
)
if self._samples is None:
raise RuntimeError("Sampling not yet conducted")
max_procs = cpu_count()
if nprocs is None:
nprocs = max_procs
else:
if nprocs > max_procs:
warnings.warn(
f"{nprocs} processors requested but only {max_procs} found."
)
nprocs = min(max_procs, nprocs)
# Create wrapped partial function to allow passing of additional args
tmp_f = self._wrap_func(func, *args, **kwargs)
# Split into even chunks
chunks = np.array_split(self._samples, int(nprocs), axis=0)
if ptqdm_available:
# Display progress bar if available
res = p_imap(tmp_f, chunks, num_cpus=nprocs)
else:
with Pool(nprocs) as pool:
res = list(pool.imap(tmp_f, chunks))
self.results = self._collect_results(res)
return self
[docs]
def evaluate_distributed(
self, func, *args, nprocs=1, servers=None, verbose=False, **kwargs
):
"""Distribute model evaluation across a cluster.
Usage Conditions:
* The provided function needs to accept a numpy array of inputs as
its first parameter
* The provided function must return a numpy array of results
Parameters
----------
func : function,
Model, or function that wraps a model, to be run in parallel
nprocs : int,
Number of processors to use for each node. Defaults to 1.
servers : list[str] or None,
IP addresses or alias for each server/node to use.
verbose : bool,
Display job execution statistics. Defaults to False.
*args : list,
Additional arguments to be passed to `func`
**kwargs : dict,
Additional keyword arguments passed to `func`
Returns
-------
self : ProblemSpec object
"""
if not pathos_available:
raise RuntimeError(
"Pathos is required to run in distributed mode. Please install"
" with `pip install pathos` or"
" `conda install pathos -c conda-forge`"
)
if verbose:
from pathos.parallel import stats
warnings.warn(
"Distributed evaluation is an untested experimental feature and"
" may not work."
)
workers = pp_Pool(nprocs, servers=servers)
# Split into even chunks
chunks = np.array_split(self._samples, int(nprocs) * len(servers), axis=0)
tmp_f = self._wrap_func(func)
res = list(workers.map(tmp_f, chunks))
self.results = self._collect_results(res)
if verbose:
print(stats(), "\n")
workers.clear()
return self
[docs]
def analyze(self, func, *args, **kwargs):
"""Analyze sampled results using given function.
Parameters
----------
func : function,
Analysis method to use. The provided function must accept the
problem specification as the first parameter, X values if needed,
Y values, and return a numpy array.
*args : list,
Additional arguments to be passed to `func`
nprocs : int,
If specified, attempts to parallelize model evaluations
**kwargs : dict,
Additional keyword arguments passed to `func`
Returns
-------
self : ProblemSpec object
"""
if self["num_vars"] == 1 or (self["groups"] and len("groups") == 1):
msg = (
"There is only a single parameter or group defined. There is "
"no point in conducting sensitivity analysis as any and all"
"effect(s) will be mapped to the single parameter/group."
)
raise ValueError(msg)
if "nprocs" in kwargs:
# Call parallel method instead
return self.analyze_parallel(func, *args, **kwargs)
if self._results is None:
raise RuntimeError("Model not yet evaluated")
if "X" in func.__code__.co_varnames:
# enforce passing of X if expected
func = partial(func, *args, X=self._samples, **kwargs)
else:
func = partial(func, *args, **kwargs)
out_cols = self.get("outputs", None)
if out_cols is None:
if len(self._results.shape) == 1:
self["outputs"] = ["Y"]
else:
num_cols = self._results.shape[1]
self["outputs"] = [f"Y{i}" for i in range(1, num_cols + 1)]
if len(self["outputs"]) > 1:
self._analysis = {}
for i, out in enumerate(self["outputs"]):
self._analysis[out] = func(self, Y=self._results[:, i])
else:
self._analysis = func(self, Y=self._results)
return self
[docs]
def analyze_parallel(self, func, *args, nprocs=None, **kwargs):
"""Analyze sampled results using the given function in parallel.
Parameters
----------
func : function,
Analysis method to use. The provided function must accept the
problem specification as the first parameter, X values if needed,
Y values, and return a numpy array.
*args : list,
Additional arguments to be passed to `func`
nprocs : int,
Number of processors to use.
Capped to the number of outputs or available processors.
**kwargs : dict,
Additional keyword arguments passed to `func`
Returns
-------
self : ProblemSpec object
"""
warnings.warn("Parallel analysis is an experimental feature and may not work.")
if self._results is None:
raise RuntimeError("Model not yet evaluated")
if "X" in func.__code__.co_varnames:
# enforce passing of X if expected
func = partial(func, *args, X=self._samples, **kwargs)
else:
func = partial(func, *args, **kwargs)
out_cols = self.get("outputs", None)
if out_cols is None:
if len(self._results.shape) == 1:
self["outputs"] = ["Y"]
else:
num_cols = self._results.shape[1]
self["outputs"] = [f"Y{i}" for i in range(1, num_cols + 1)]
# Cap number of processors used
Yn = len(self["outputs"])
if Yn == 1:
# Only single output, cannot parallelize
warnings.warn(
f"Analysis was not parallelized: {nprocs} processors requested"
f" for 1 output."
)
res = func(self, Y=self._results)
else:
max_procs = cpu_count()
if nprocs is None:
nprocs = max_procs
else:
nprocs = min(Yn, nprocs, max_procs)
if ptqdm_available:
# Display progress bar if available
res = p_imap(
lambda y: func(self, Y=y),
[self._results[:, i] for i in range(Yn)],
num_cpus=nprocs,
)
else:
with Pool(nprocs) as pool:
res = list(
pool.imap(
lambda y: func(self, Y=y),
[self._results[:, i] for i in range(Yn)],
)
)
# Assign by output name if more than 1 output, otherwise
# attach directly
if Yn > 1:
self._analysis = {}
for out, Si in zip(self["outputs"], list(res)):
self._analysis[out] = Si
else:
self._analysis = res
return self
[docs]
def to_df(self):
"""Convert results to Pandas DataFrame."""
an_res = self._analysis
if isinstance(an_res, ResultDict):
return an_res.to_df()
elif isinstance(an_res, dict):
# case where analysis result is a dict of ResultDicts
return [an.to_df() for an in list(an_res.values())]
raise RuntimeError("Analysis not yet conducted")
[docs]
def plot(self, **kwargs):
"""Plot results as a bar chart.
Returns
-------
axes : matplotlib axes object
"""
if self._analysis is None:
raise RuntimeError("Analysis not yet conducted")
num_rows = len(self["outputs"])
if num_rows == 1:
return self._analysis.plot(**kwargs)
try:
plt
except NameError:
import matplotlib.pyplot as plt
num_cols = 1
fk = list(self._analysis.keys())[0]
if isinstance(self._analysis[fk].to_df(), (list, tuple)):
# have to divide by 2 to account for CI columns
num_cols = len(self._analysis[fk]) // 2
p_width = max(num_cols * 3, 5)
p_height = max(num_rows * 3, 6)
_, axes = plt.subplots(
num_rows, num_cols, sharey=True, figsize=(p_width, p_height)
)
for res, ax in zip(self._analysis, axes):
self._analysis[res].plot(ax=ax, **kwargs)
try:
ax[0].set_title(res)
except TypeError:
ax.set_title(res)
plt.tight_layout()
return axes
[docs]
def heatmap(
self, metric: str = None, index: str = None, title: str = None, ax=None
):
"""Plot results as a heatmap.
Parameters
----------
metric : str or None, name of output to analyze (display all if `None`)
index : str or None, name of index to plot, dependent on what
analysis was conducted (ST, S1, etc; displays all if `None`)
title : str, title of plot to use (defaults to the same as `metric`)
ax : axes object, matplotlib axes object to use for plot.
Creates a new figure if not provided.
Returns
-------
ax : matplotlib axes object
"""
from SALib.plotting.heatmap import heatmap # type: ignore
return heatmap(self, metric, index, title, ax)
def _wrap_func(self, func, *args, **kwargs):
# Create wrapped partial function to allow passing of additional args
tmp_f = func
if (len(args) > 0) or (len(kwargs) > 0):
tmp_f = lambda x: func(x, *args, **kwargs) # noqa
return tmp_f
def _collect_results(self, res):
res_shape = res[0].shape
if len(res_shape) > 1:
res_shape = (len(self._samples), *res_shape[1:])
else:
res_shape = len(self._samples)
final_res = np.empty(res_shape)
# Collect results
# Cannot enumerate over this as the length
# of individual results may vary
i = 0
for r in res:
r_len = len(r)
final_res[i : i + r_len] = r
i += r_len
return final_res
def _method_creator(self, func, method):
"""Generate convenience methods for specified `method`."""
@wraps(func)
def modfunc(self, *args, **kwargs):
return getattr(self, method)(func, *args, **kwargs)
return modfunc
def _add_samplers(self):
"""Dynamically add available SALib samplers as ProblemSpec methods."""
for sampler in avail_approaches(samplers):
func = getattr(
importlib.import_module("SALib.sample.{}".format(sampler)), "sample"
)
method_name = "sample_{}".format(sampler.replace("_sampler", ""))
self.__setattr__(
method_name, MethodType(self._method_creator(func, "sample"), self)
)
def _add_analyzers(self):
"""Dynamically add available SALib analyzers as ProblemSpec methods."""
for analyzer in avail_approaches(analyzers):
func = getattr(
importlib.import_module("SALib.analyze.{}".format(analyzer)), "analyze"
)
method_name = "analyze_{}".format(analyzer.replace("_analyzer", ""))
self.__setattr__(
method_name, MethodType(self._method_creator(func, "analyze"), self)
)
def _repr_pretty_(self, p, cycle):
p.text(str(self) if not cycle else "...")
def __str__(self):
rep = ""
if self._samples is not None:
arr_shape = self._samples.shape
if len(arr_shape) == 1:
arr_shape = (arr_shape[0], 1)
nr, nx = arr_shape
rep += (
"Samples:\n" f"\t{nx} parameters: {self['names']}\n" f"\t{nr} samples\n"
)
if self._results is not None:
arr_shape = self._results.shape
if len(arr_shape) == 1:
arr_shape = (arr_shape[0], 1)
nr, ny = arr_shape
rep += (
"Outputs:\n"
f"\t{ny} outputs: {self['outputs']}\n"
f"\t{nr} evaluations\n"
)
if self._analysis is not None:
rep += "Analysis:\n"
an_res = self._analysis
allowed_types = (list, tuple)
if isinstance(an_res, ResultDict):
an_res = an_res.to_df()
if not isinstance(an_res, allowed_types):
rep += f"{an_res}\n"
else:
for df in an_res:
rep += f"{df}\n"
elif isinstance(an_res, dict):
for res_name in an_res:
rep += "{}:\n".format(res_name)
dfs = an_res[res_name].to_df()
if isinstance(dfs, allowed_types):
for df in dfs:
rep += f"{df}:\n"
else:
rep += f"{dfs}:\n"
if len(rep) == 0:
rep = (
"ProblemSpec does not currently contain any samples, "
"evaluations or results."
)
return rep
def _check_spec_attributes(spec: ProblemSpec):
assert "names" in spec, "Names not defined"
assert "bounds" in spec, "Bounds not defined"
assert len(spec["bounds"]) == len(
spec["names"]
), f"""Number of bounds do not match number of names
Number of names:
{len(spec['names'])} | {spec['names']}
----------------
Number of bounds: {len(spec['bounds'])}
"""
