def test_relations(index_dependent, link_clp):
model = deepcopy(suite.model)
model.dataset_group_models["default"].link_clp = link_clp
model.megacomplex["m1"].is_index_dependent = index_dependent
model.clp_relations.append(
Relation.from_dict({
"source": "s1",
"target": "s2",
"parameter": "3"
}))
parameters = ParameterGroup.from_list([11e-4, 22e-5, 2])
print("link_clp", link_clp, "index_dependent", index_dependent)
dataset = simulate(
suite.sim_model,
"dataset1",
parameters,
{
"global": suite.global_axis,
"model": suite.model_axis
},
)
scheme = Scheme(model=model,
parameters=parameters,
data={"dataset1": dataset})
optimization_group = OptimizationGroup(
scheme,
model.get_dataset_groups()["default"])
if index_dependent:
reduced_matrix = (optimization_group.reduced_matrices[0]
if link_clp else
optimization_group.reduced_matrices["dataset1"][0])
else:
reduced_matrix = optimization_group.reduced_matrices["dataset1"]
matrix = (optimization_group.matrices["dataset1"][0]
if index_dependent else optimization_group.matrices["dataset1"])
result_data = optimization_group.create_result_data()
print(result_data)
clps = result_data["dataset1"].clp
assert "s2" not in reduced_matrix.clp_labels
assert "s2" in clps.coords["clp_label"]
assert clps.sel(clp_label="s2") == clps.sel(clp_label="s1") * 2
assert "s2" in matrix.clp_labels
def test_full_model_problem():
dataset = simulate(FullModel.model, "dataset1", FullModel.parameters,
FullModel.coordinates)
scheme = Scheme(model=FullModel.model,
parameters=FullModel.parameters,
data={"dataset1": dataset})
optimization_group = OptimizationGroup(
scheme,
FullModel.model.get_dataset_groups()["default"])
result = optimization_group.create_result_data()["dataset1"]
assert "global_matrix" in result
assert "global_clp_label" in result
clp = result.clp
assert clp.shape == (4, 4)
print(np.diagonal(clp))
assert all(np.isclose(1.0, c) for c in np.diagonal(clp))
def test_single_dataset():
model = SimpleTestModel.from_dict({
"megacomplex": {
"m1": {
"is_index_dependent": False
}
},
"dataset_groups": {
"default": {
"link_clp": True
}
},
"dataset": {
"dataset1": {
"megacomplex": ["m1"],
},
},
})
print(model.validate())
assert model.valid()
parameters = ParameterGroup.from_list([1, 10])
print(model.validate(parameters))
assert model.valid(parameters)
global_axis = [1, 2, 3]
model_axis = [5, 7, 9, 12]
data = {
"dataset1":
xr.DataArray(np.ones((3, 4)),
coords=[("global", global_axis),
("model", model_axis)]).to_dataset(name="data")
}
scheme = Scheme(model, parameters, data)
optimization_group = OptimizationGroup(
scheme,
model.get_dataset_groups()["default"])
bag = optimization_group._calculator.bag
datasets = optimization_group._calculator.groups
assert len(datasets) == 1
assert len(bag) == 3
assert all(p.data.size == 4 for p in bag)
assert all(p.dataset_models[0].label == "dataset1" for p in bag)
assert all(
all(p.dataset_models[0].axis["model"] == model_axis) for p in bag)
assert all(
all(p.dataset_models[0].axis["global"] == global_axis) for p in bag)
assert [p.dataset_models[0].indices["global"] for p in bag] == [0, 1, 2]
def test_optimization_group_result_data(optimization_group: OptimizationGroup):
data = optimization_group.create_result_data()
label = "dataset1"
assert label in data
dataset = data[label]
dataset_model = optimization_group.dataset_models[label]
assert "clp_label" in dataset.coords
assert np.array_equal(dataset.clp_label, ["s1", "s2", "s3", "s4"])
assert dataset_model.get_global_dimension() in dataset.coords
assert np.array_equal(dataset.coords[dataset_model.get_global_dimension()],
suite.global_axis)
assert dataset_model.get_model_dimension() in dataset.coords
assert np.array_equal(dataset.coords[dataset_model.get_model_dimension()],
suite.model_axis)
assert "matrix" in dataset
matrix = dataset.matrix
if optimization_group.model.is_index_dependent:
assert len(matrix.shape) == 3
assert matrix.shape[0] == suite.global_axis.size
assert matrix.shape[1] == suite.model_axis.size
assert matrix.shape[2] == 4
else:
assert len(matrix.shape) == 2
assert matrix.shape[0] == suite.model_axis.size
assert matrix.shape[1] == 4
assert "clp" in dataset
clp = dataset.clp
assert len(clp.shape) == 2
assert clp.shape[0] == suite.global_axis.size
assert clp.shape[1] == 4
assert "weighted_residual" in dataset
assert dataset.data.shape == dataset.weighted_residual.shape
assert "residual" in dataset
assert dataset.data.shape == dataset.residual.shape
assert "residual_singular_values" in dataset
assert "weighted_residual_singular_values" in dataset
def optimization_group(request) -> OptimizationGroup:
model = suite.model
model.megacomplex["m1"].is_index_dependent = request.param[1]
model.is_index_dependent = request.param[1]
model.dataset_group_models["default"].link_clp = request.param[0]
dataset = simulate(
suite.sim_model,
"dataset1",
suite.wanted_parameters,
{
"global": suite.global_axis,
"model": suite.model_axis
},
)
scheme = Scheme(model=model,
parameters=suite.initial_parameters,
data={"dataset1": dataset})
return OptimizationGroup(scheme, model.get_dataset_groups()["default"])
def test_penalties(index_dependent, link_clp):
model = deepcopy(suite.model)
model.dataset_group_models["default"].link_clp = link_clp
model.megacomplex["m1"].is_index_dependent = index_dependent
model.clp_area_penalties.append(
EqualAreaPenalty.from_dict({
"source": "s1",
"source_intervals": [(1, 20)],
"target": "s2",
"target_intervals": [(20, 45)],
"parameter": "3",
"weight": 10,
}))
parameters = ParameterGroup.from_list([11e-4, 22e-5, 2])
global_axis = np.arange(50)
print(f"{link_clp=}\n{index_dependent=}")
dataset = simulate(
suite.sim_model,
"dataset1",
parameters,
{
"global": global_axis,
"model": suite.model_axis
},
)
scheme = Scheme(model=model,
parameters=parameters,
data={"dataset1": dataset})
optimization_group = OptimizationGroup(
scheme,
model.get_dataset_groups()["default"])
assert isinstance(optimization_group.additional_penalty, np.ndarray)
assert optimization_group.additional_penalty.size == 1
assert optimization_group.additional_penalty[0] != 0
assert isinstance(optimization_group.full_penalty, np.ndarray)
assert (optimization_group.full_penalty.size == (suite.model_axis.size *
global_axis.size) +
optimization_group.additional_penalty.size)
def optimize(scheme: Scheme,
verbose: bool = True,
raise_exception: bool = False) -> Result:
optimization_groups = [
OptimizationGroup(scheme, group)
for group in scheme.model.get_dataset_groups().values()
]
(
free_parameter_labels,
initial_parameter,
lower_bounds,
upper_bounds,
) = scheme.parameters.get_label_value_and_bounds_arrays(
exclude_non_vary=True)
if scheme.optimization_method not in SUPPORTED_METHODS:
raise ValueError(
f"Unsupported optimization method {scheme.optimization_method}. "
f"Supported methods are '{list(SUPPORTED_METHODS.keys())}'")
method = SUPPORTED_METHODS[scheme.optimization_method]
nfev = scheme.maximum_number_function_evaluations
ftol = scheme.ftol
gtol = scheme.gtol
xtol = scheme.xtol
verbose = 2 if verbose else 0
termination_reason = ""
parameter_history = ParameterHistory()
parameter_history.append(scheme.parameters)
try:
ls_result = least_squares(
_calculate_penalty,
initial_parameter,
bounds=(lower_bounds, upper_bounds),
method=method,
max_nfev=nfev,
verbose=verbose,
ftol=ftol,
gtol=gtol,
xtol=xtol,
kwargs={
"free_parameter_labels": free_parameter_labels,
"optimization_groups": optimization_groups,
"parameter_history": parameter_history,
},
)
termination_reason = ls_result.message
except Exception as e:
if raise_exception:
raise e
warn(f"Optimization failed:\n\n{e}")
termination_reason = str(e)
ls_result = None
return _create_result(
scheme,
optimization_groups,
ls_result,
free_parameter_labels,
termination_reason,
parameter_history,
)
def setup_optimization_group(scheme):
return OptimizationGroup(scheme, scheme.model.get_dataset_groups()["default"])
def test_multi_dataset_no_overlap():
model = SimpleTestModel.from_dict({
"megacomplex": {
"m1": {
"is_index_dependent": False
}
},
"dataset_groups": {
"default": {
"link_clp": True
}
},
"dataset": {
"dataset1": {
"megacomplex": ["m1"],
},
"dataset2": {
"megacomplex": ["m1"],
},
},
})
model.grouped = lambda: True
print(model.validate())
assert model.valid()
assert model.grouped()
parameters = ParameterGroup.from_list([1, 10])
print(model.validate(parameters))
assert model.valid(parameters)
global_axis_1 = [1, 2, 3]
model_axis_1 = [5, 7]
global_axis_2 = [4, 5, 6]
model_axis_2 = [5, 7, 9]
data = {
"dataset1":
xr.DataArray(np.ones((3, 2)),
coords=[("global", global_axis_1),
("model", model_axis_1)]).to_dataset(name="data"),
"dataset2":
xr.DataArray(np.ones((3, 3)),
coords=[("global", global_axis_2),
("model", model_axis_2)]).to_dataset(name="data"),
}
scheme = Scheme(model, parameters, data)
optimization_group = OptimizationGroup(
scheme,
model.get_dataset_groups()["default"])
bag = list(optimization_group._calculator.bag)
assert len(optimization_group._calculator.groups) == 2
assert len(bag) == 6
assert all(p.data.size == 2 for p in bag[:3])
assert all(p.dataset_models[0].label == "dataset1" for p in bag[:3])
assert all(
all(p.dataset_models[0].axis["model"] == model_axis_1)
for p in bag[:3])
assert all(
all(p.dataset_models[0].axis["global"] == global_axis_1)
for p in bag[:3])
assert [p.dataset_models[0].indices["global"]
for p in bag[:3]] == [0, 1, 2]
assert all(p.data.size == 3 for p in bag[3:])
assert all(p.dataset_models[0].label == "dataset2" for p in bag[3:])
assert all(
all(p.dataset_models[0].axis["model"] == model_axis_2)
for p in bag[3:])
assert all(
all(p.dataset_models[0].axis["global"] == global_axis_2)
for p in bag[3:])
assert [p.dataset_models[0].indices["global"]
for p in bag[3:]] == [0, 1, 2]
def test_multi_dataset_overlap():
model = SimpleTestModel.from_dict({
"megacomplex": {
"m1": {
"is_index_dependent": False
}
},
"dataset_groups": {
"default": {
"link_clp": True
}
},
"dataset": {
"dataset1": {
"megacomplex": ["m1"],
},
"dataset2": {
"megacomplex": ["m1"],
},
},
})
model.grouped = lambda: True
print(model.validate())
assert model.valid()
assert model.grouped()
parameters = ParameterGroup.from_list([1, 10])
print(model.validate(parameters))
assert model.valid(parameters)
global_axis_1 = [1, 2, 3, 5]
model_axis_1 = [5, 7]
global_axis_2 = [0, 1.4, 2.4, 3.4, 9]
model_axis_2 = [5, 7, 9, 12]
data = {
"dataset1":
xr.DataArray(np.ones((4, 2)),
coords=[("global", global_axis_1),
("model", model_axis_1)]).to_dataset(name="data"),
"dataset2":
xr.DataArray(np.ones((5, 4)),
coords=[("global", global_axis_2),
("model", model_axis_2)]).to_dataset(name="data"),
}
scheme = Scheme(model, parameters, data, clp_link_tolerance=5e-1)
optimization_group = OptimizationGroup(
scheme,
model.get_dataset_groups()["default"])
bag = list(optimization_group._calculator.bag)
assert len(optimization_group._calculator.groups) == 3
assert "dataset1dataset2" in optimization_group._calculator.groups
assert optimization_group._calculator.groups["dataset1dataset2"] == [
"dataset1", "dataset2"
]
assert len(bag) == 6
assert all(p.data.size == 4 for p in bag[:1])
assert all(p.dataset_models[0].label == "dataset1" for p in bag[1:5])
assert all(
all(p.dataset_models[0].axis["model"] == model_axis_1)
for p in bag[1:5])
assert all(
all(p.dataset_models[0].axis["global"] == global_axis_1)
for p in bag[1:5])
assert [p.dataset_models[0].indices["global"]
for p in bag[1:5]] == [0, 1, 2, 3]
assert all(p.data.size == 6 for p in bag[1:4])
assert all(p.dataset_models[1].label == "dataset2" for p in bag[1:4])
assert all(
all(p.dataset_models[1].axis["model"] == model_axis_2)
for p in bag[1:4])
assert all(
all(p.dataset_models[1].axis["global"] == global_axis_2)
for p in bag[1:4])
assert [p.dataset_models[1].indices["global"]
for p in bag[1:4]] == [1, 2, 3]
assert all(p.data.size == 4 for p in bag[5:])
assert bag[4].dataset_models[0].label == "dataset1"
assert bag[5].dataset_models[0].label == "dataset2"
assert np.array_equal(bag[4].dataset_models[0].axis["model"], model_axis_1)
assert np.array_equal(bag[5].dataset_models[0].axis["model"], model_axis_2)
assert [p.dataset_models[0].indices["global"]
for p in bag[1:4]] == [0, 1, 2]
def test_prepare_data():
model_dict = {
"megacomplex": {
"m1": {
"is_index_dependent": False
}
},
"dataset": {
"dataset1": {
"megacomplex": ["m1"],
},
},
"weights": [
{
"datasets": ["dataset1"],
"global_interval": (np.inf, 200),
"model_interval": (4, 8),
"value": 0.5,
},
],
}
model = SimpleTestModel.from_dict(model_dict)
print(model.validate())
assert model.valid()
parameters = ParameterGroup.from_list([])
global_axis = np.asarray(range(50, 300))
model_axis = np.asarray(range(15))
dataset = xr.DataArray(
np.ones((global_axis.size, model_axis.size)),
coords={
"global": global_axis,
"model": model_axis
},
dims=("global", "model"),
)
scheme = Scheme(model, parameters, {"dataset1": dataset})
optimization_group = OptimizationGroup(
scheme,
model.get_dataset_groups()["default"])
data = optimization_group.data["dataset1"]
print(data)
assert "data" in data
assert "weight" in data
assert data.data.shape == (model_axis.size, global_axis.size)
assert data.data.shape == data.weight.shape
assert np.all(
data.weight.sel({
"global": slice(0, 200),
"model": slice(4, 8)
}).values == 0.5)
assert np.all(data.weight.sel(model=slice(0, 3)).values == 1)
model_dict["weights"].append({
"datasets": ["dataset1"],
"value": 0.2,
})
model = SimpleTestModel.from_dict(model_dict)
print(model.validate())
assert model.valid()
scheme = Scheme(model, parameters, {"dataset1": dataset})
optimization_group = OptimizationGroup(
scheme,
model.get_dataset_groups()["default"])
data = optimization_group.data["dataset1"]
assert np.all(
data.weight.sel({
"global": slice(0, 200),
"model": slice(4, 8)
}).values == 0.5 * 0.2)
assert np.all(data.weight.sel(model=slice(0, 3)).values == 0.2)
with pytest.warns(
UserWarning,
match="Ignoring model weight for dataset 'dataset1'"
" because weight is already supplied by dataset.",
):
OptimizationGroup(Scheme(model, parameters, {"dataset1": data}),
model.get_dataset_groups()["default"])