diff --git a/corrai/base/distribution.py b/corrai/base/distribution.py new file mode 100644 index 0000000..2c43382 --- /dev/null +++ b/corrai/base/distribution.py @@ -0,0 +1,87 @@ +from dataclasses import dataclass, field + +import numpy as np +from scipy import stats + +DISTRIBUTIONS = ["normal", "truncnormal", "uniform", "halfcauchy", "lognormal"] + +REQUIRED_PARAMS = { + "normal": {"mean", "std"}, + "truncnormal": {"mean", "std", "low", "high"}, + "uniform": {"low", "high"}, + "halfcauchy": {"loc", "scale"}, + "lognormal": {"mean", "sigma"}, +} + + +@dataclass +class Distribution: + """ + Probability distribution assigned to a `Parameter`, used to draw random + values for uncertainty propagation (see `corrai.sampling.MonteCarloSampler`). + + Parameters + ---------- + dist : str + Distribution family. One of `DISTRIBUTIONS`: + `"normal"`, `"truncnormal"`, `"uniform"`, `"halfcauchy"`, `"lognormal"`. + params : dict + Distribution parameters. Required keys depend on `dist`: + + - `"normal"`: `mean`, `std` + - `"truncnormal"`: `mean`, `std`, `low`, `high` + - `"uniform"`: `low`, `high` + - `"halfcauchy"`: `loc`, `scale` + - `"lognormal"`: `mean`, `sigma` (parameters of the underlying normal) + + Examples + -------- + >>> Distribution("normal", {"mean": 0.036, "std": 0.002}) + >>> Distribution("truncnormal", {"mean": 0.036, "std": 0.002, "low": 0.03, "high": 0.04}) + >>> Distribution("uniform", {"low": 0.03, "high": 0.04}) + >>> Distribution("halfcauchy", {"loc": 0, "scale": 1}) + """ + + dist: str + params: dict = field(default_factory=dict) + + def __post_init__(self): + if self.dist not in DISTRIBUTIONS: + raise ValueError( + f"Invalid distribution: {self.dist!r}. Must be one of {DISTRIBUTIONS}." + ) + + required = REQUIRED_PARAMS[self.dist] + missing = required - self.params.keys() + if missing: + raise ValueError( + f"Missing parameters {sorted(missing)} for distribution {self.dist!r}. " + f"Required: {sorted(required)}." + ) + + def _frozen(self): + p = self.params + if self.dist == "normal": + return stats.norm(loc=p["mean"], scale=p["std"]) + if self.dist == "truncnormal": + a = (p["low"] - p["mean"]) / p["std"] + b = (p["high"] - p["mean"]) / p["std"] + return stats.truncnorm(a, b, loc=p["mean"], scale=p["std"]) + if self.dist == "uniform": + return stats.uniform(loc=p["low"], scale=p["high"] - p["low"]) + if self.dist == "halfcauchy": + return stats.halfcauchy(loc=p["loc"], scale=p["scale"]) + if self.dist == "lognormal": + return stats.lognorm(s=p["sigma"], scale=np.exp(p["mean"])) + raise ValueError(f"Invalid distribution: {self.dist!r}") + + def rvs(self, size, random_state=None): + return self._frozen().rvs(size=size, random_state=random_state) + + def quantile_range(self, low: float = 0.005, high: float = 0.995): + """ + Return a display-friendly (low, high) range, used as a fallback + interval for plotting when a `Parameter` has no explicit `interval`. + """ + lo, hi = self._frozen().ppf([low, high]) + return float(lo), float(hi) diff --git a/corrai/base/parameter.py b/corrai/base/parameter.py index 8decf01..95ba3a1 100644 --- a/corrai/base/parameter.py +++ b/corrai/base/parameter.py @@ -1,5 +1,7 @@ from dataclasses import dataclass +from corrai.base.distribution import Distribution + TYPES = ["Integer", "Real", "Choice", "Binary"] RELABS = ["Absolute", "Relative"] @@ -16,6 +18,7 @@ class Parameter: tuple[int | float, int | float] | list[tuple[int | float, int | float]] | None ) = None model_property: str | tuple[str, ...] = None + distribution: Distribution | None = None """ A parameter definition for models. Can Affect a single model property or a list @@ -64,12 +67,23 @@ class Parameter: min_max_interval : tuple of int or float, optional Optional min and max bounds used for some checking operations. + distribution : Distribution, optional + Probability distribution used to draw random values for this + parameter, e.g. with `corrai.sampling.MonteCarloSampler` for + uncertainty propagation. Only valid for `ptype="Real"`. If + `interval` is not provided, at least `distribution` or `values` + must be. `interval` may still be provided alongside `distribution` + as a purely informative/plotting bound (it does not constrain the + draws). + Raises ------ ValueError - If both `interval` and `values` are specified, or if neither is specified. + If both `interval` and `values` are specified, or if none of + `interval`, `values`, or `distribution` is specified. If `init_value` is outside the specified domain. If `ptype` or `relabs` are not in the allowed sets. + If `distribution` is specified with a `ptype` other than `"Real"`. Examples -------- @@ -93,17 +107,38 @@ class Parameter: ... ptype="Choice", ... init_value="TARP" ... ) + + >>> # Example using a probability distribution for uncertainty propagation + >>> from corrai.base.distribution import Distribution + >>> p = Parameter( + ... name="Conductivity", + ... model_property="building.wall.insulation.conductivity", + ... distribution=Distribution( + ... "truncnormal", {"mean": 0.036, "std": 0.002, "low": 0.03, "high": 0.04} + ... ), + ... ) """ def __post_init__(self): if self.interval is not None and self.values is not None: raise ValueError("Only one of 'interval' or 'values' may be specified.") - if self.interval is None and self.values is None and self.ptype != "Binary": + if ( + self.interval is None + and self.values is None + and self.distribution is None + and self.ptype != "Binary" + ): raise ValueError("One of 'interval' or 'values' must be specified.") if self.ptype not in TYPES: raise ValueError(f"Invalid type: {self.ptype!r}. Must be one of {TYPES}.") + if self.distribution is not None and self.ptype != "Real": + raise ValueError( + f"'distribution' can only be used with ptype='Real', " + f"got ptype={self.ptype!r}." + ) + if isinstance(self.relabs, str) and self.relabs not in RELABS: raise ValueError( f"Invalid relabs: {self.relabs!r}. " diff --git a/corrai/sampling.py b/corrai/sampling.py index 2029e3c..115d030 100644 --- a/corrai/sampling.py +++ b/corrai/sampling.py @@ -218,7 +218,14 @@ def get_parameters_intervals(self): If parameters are not of type 'Real'. """ if all(param.ptype == "Real" for param in self.parameters): - return np.array([param.interval for param in self.parameters]) + return np.array( + [ + param.interval + if param.interval is not None + else param.distribution.quantile_range() + for param in self.parameters + ] + ) elif any(param.ptype == "Integer" for param in self.parameters): raise NotImplementedError( "get_param_interval is not yet implemented for integer parameters" @@ -1569,3 +1576,69 @@ def add_sample( sample_is_dimless=False, simulation_kwargs=simulation_kwargs, ) + + +class MonteCarloSampler(RealSampler): + """ + Monte Carlo sampler for uncertainty propagation. + + Draws each parameter independently from its `Parameter.distribution` + and runs simulations for the resulting samples. Used to propagate + parameter uncertainty (given as probability distributions) through the + model, rather than to explore a deterministic design of experiments. + + Parameters + ---------- + parameters : list of Parameter + Real-valued parameters, each with a `distribution` set. + model : Model + Model to simulate. + simulation_options : dict, optional + Options for simulation. + + Raises + ------ + ValueError + If any parameter does not define a `distribution`. + + Methods + ------- + add_sample(n, seed=None, simulate=True, n_cpu=1, simulation_kwargs=None) + Draw `n` samples from the parameters' distributions. + """ + + def __init__( + self, + parameters: list[Parameter], + model: Model, + simulation_options: dict = None, + ): + super().__init__(parameters, model, simulation_options) + + missing = [par.name for par in parameters if par.distribution is None] + if missing: + raise ValueError( + f"All parameters must define a `distribution`. Missing for: {missing}" + ) + + def add_sample( + self, + n: int, + seed: int = None, + simulate: bool = True, + n_cpu: int = 1, + simulation_kwargs: dict = None, + ): + rng = np.random.default_rng(seed) + columns = [ + param.distribution.rvs(size=n, random_state=rng) + for param in self.parameters + ] + new_sample = np.column_stack(columns) + self._post_draw_sample( + new_sample, + simulate, + n_cpu, + sample_is_dimless=False, + simulation_kwargs=simulation_kwargs, + ) diff --git a/tests/base/test_distribution.py b/tests/base/test_distribution.py new file mode 100644 index 0000000..b92e602 --- /dev/null +++ b/tests/base/test_distribution.py @@ -0,0 +1,57 @@ +import numpy as np +import pytest + +from corrai.base.distribution import Distribution + + +class TestDistribution: + def test_invalid_dist_name(self): + with pytest.raises(ValueError, match="Invalid distribution"): + Distribution("not_a_dist", {}) + + def test_missing_params(self): + with pytest.raises(ValueError, match="Missing parameters"): + Distribution("normal", {"mean": 0}) + + def test_normal_rvs(self): + dist = Distribution("normal", {"mean": 10, "std": 2}) + samples = dist.rvs(size=5000, random_state=0) + assert samples.shape == (5000,) + assert np.isclose(samples.mean(), 10, atol=0.2) + assert np.isclose(samples.std(), 2, atol=0.2) + + def test_truncnormal_rvs_within_bounds(self): + dist = Distribution( + "truncnormal", {"mean": 0.036, "std": 0.01, "low": 0.03, "high": 0.04} + ) + samples = dist.rvs(size=2000, random_state=0) + assert samples.min() >= 0.03 + assert samples.max() <= 0.04 + + def test_uniform_rvs_within_bounds(self): + dist = Distribution("uniform", {"low": 2, "high": 5}) + samples = dist.rvs(size=2000, random_state=0) + assert samples.min() >= 2 + assert samples.max() <= 5 + + def test_halfcauchy_rvs_non_negative(self): + dist = Distribution("halfcauchy", {"loc": 0, "scale": 1}) + samples = dist.rvs(size=2000, random_state=0) + assert samples.min() >= 0 + + def test_lognormal_rvs_positive(self): + dist = Distribution("lognormal", {"mean": 0, "sigma": 0.5}) + samples = dist.rvs(size=2000, random_state=0) + assert samples.min() > 0 + + def test_quantile_range(self): + dist = Distribution("uniform", {"low": 0, "high": 10}) + low, high = dist.quantile_range() + assert low == pytest.approx(0.05, abs=1e-6) + assert high == pytest.approx(9.95, abs=1e-6) + + def test_rvs_reproducible_with_seed(self): + dist = Distribution("normal", {"mean": 0, "std": 1}) + s1 = dist.rvs(size=100, random_state=np.random.default_rng(42)) + s2 = dist.rvs(size=100, random_state=np.random.default_rng(42)) + np.testing.assert_allclose(s1, s2) diff --git a/tests/base/test_parameter.py b/tests/base/test_parameter.py index ae5bc08..da48190 100644 --- a/tests/base/test_parameter.py +++ b/tests/base/test_parameter.py @@ -1,5 +1,6 @@ import pytest from corrai.base.parameter import Parameter +from corrai.base.distribution import Distribution VALID_INTERVAL = (0, 10) @@ -71,3 +72,32 @@ def test_init_value_not_in_values(self): values=("A", "B", "C"), init_value="D", ) + + def test_parameter_with_distribution_only(self): + param = Parameter( + name="x", + model_property="m.x", + distribution=Distribution("normal", {"mean": 0, "std": 1}), + ) + assert param.interval is None + assert param.distribution.dist == "normal" + + def test_parameter_with_distribution_and_interval(self): + param = Parameter( + name="x", + model_property="m.x", + interval=(0, 1), + distribution=Distribution("uniform", {"low": 0, "high": 1}), + ) + assert param.interval == (0, 1) + assert param.distribution is not None + + def test_distribution_requires_real_ptype(self): + with pytest.raises(ValueError, match="distribution.*ptype='Real'"): + Parameter( + name="bad", + model_property="m.bad", + values=("A", "B"), + ptype="Choice", + distribution=Distribution("normal", {"mean": 0, "std": 1}), + ) diff --git a/tests/test_sampling.py b/tests/test_sampling.py index 9598369..f95b8f0 100644 --- a/tests/test_sampling.py +++ b/tests/test_sampling.py @@ -3,10 +3,12 @@ import pandas as pd from corrai.base.parameter import Parameter +from corrai.base.distribution import Distribution from corrai.sampling import ( LHSSampler, MorrisSampler, SobolSampler, + MonteCarloSampler, Sample, ) @@ -43,6 +45,27 @@ Parameter("par_x3", (-3.14159265359, 3.14159265359), model_property="x3"), ] +MONTECARLO_PARAMS = [ + Parameter( + "param_1", + model_property="prop_1", + distribution=Distribution("normal", {"mean": 5, "std": 1}), + ), + Parameter( + "param_2", + model_property="prop_2", + relabs="Relative", + distribution=Distribution("uniform", {"low": 0.8, "high": 1.2}), + ), + Parameter( + "param_3", + model_property="prop_3", + distribution=Distribution( + "truncnormal", {"mean": 50, "std": 20, "low": 0, "high": 100} + ), + ), +] + class TestSample: def test_sample_methods(self): @@ -560,3 +583,47 @@ def test_sobol_sampler(self): ] ), ) + + def test_montecarlo_sampler(self): + sampler = MonteCarloSampler( + parameters=MONTECARLO_PARAMS, + model=PymodelDynamic(), + simulation_options=SIMULATION_OPTIONS, + ) + sampler.add_sample(2000, seed=42, simulate=False) + + assert sampler.values.shape == (2000, 3) + np.testing.assert_allclose(sampler.values["param_1"].mean(), 5, atol=0.2) + np.testing.assert_allclose(sampler.values["param_1"].std(), 1, atol=0.2) + assert sampler.values["param_2"].min() >= 0.8 + assert sampler.values["param_2"].max() <= 1.2 + assert sampler.values["param_3"].min() >= 0 + assert sampler.values["param_3"].max() <= 100 + + sampler.simulate_at(0) + assert not sampler.results.iloc[0].empty + + def test_montecarlo_sampler_reproducible_with_seed(self): + sampler_1 = MonteCarloSampler( + parameters=MONTECARLO_PARAMS, + model=PymodelDynamic(), + simulation_options=SIMULATION_OPTIONS, + ) + sampler_1.add_sample(10, seed=42, simulate=False) + + sampler_2 = MonteCarloSampler( + parameters=MONTECARLO_PARAMS, + model=PymodelDynamic(), + simulation_options=SIMULATION_OPTIONS, + ) + sampler_2.add_sample(10, seed=42, simulate=False) + + pd.testing.assert_frame_equal(sampler_1.values, sampler_2.values) + + def test_montecarlo_sampler_missing_distribution_raises(self): + with pytest.raises(ValueError, match="must define a `distribution`"): + MonteCarloSampler( + parameters=REAL_PARAM, + model=PymodelDynamic(), + simulation_options=SIMULATION_OPTIONS, + ) diff --git a/tests/test_store.py b/tests/test_store.py index 0d05b46..d9d9c54 100644 --- a/tests/test_store.py +++ b/tests/test_store.py @@ -89,8 +89,14 @@ def test_name_collision_is_handled(self, tmp_path): bundle = tmp_path / "bundle" bundle.mkdir() + opts = {"file_a": dir_a / "data.csv", "file_b": dir_b / "data.csv"} + packed = _pack_simulation_options(opts, bundle) + files = list((bundle / "simulation_files").iterdir()) assert len(files) == 2 + assert {f.name for f in files} == {"data.csv", "data_1.csv"} + assert packed["file_a"]["__bundle_file__"] == "simulation_files/data.csv" + assert packed["file_b"]["__bundle_file__"] == "simulation_files/data_1.csv" def test_nonexistent_path_string_not_treated_as_file(self, tmp_path): bundle = tmp_path / "bundle"