From 2f6d90991bf826aad6f87fd556bb5a2d0e642b65 Mon Sep 17 00:00:00 2001 From: Marc Planelles Date: Thu, 17 Oct 2024 18:09:47 +0200 Subject: [PATCH] feat: Add support for non_parametric optimization --- src/ansys/simai/core/api/optimization.py | 8 +- src/ansys/simai/core/data/optimizations.py | 160 ++++++++++++++---- tests/test_optimizations.py | 185 ++++++++++++++++++++- 3 files changed, 312 insertions(+), 41 deletions(-) diff --git a/src/ansys/simai/core/api/optimization.py b/src/ansys/simai/core/api/optimization.py index 0b067d66..0e88403c 100644 --- a/src/ansys/simai/core/api/optimization.py +++ b/src/ansys/simai/core/api/optimization.py @@ -29,12 +29,10 @@ class OptimizationClientMixin(ApiClientMixin): def define_optimization(self, workspace_id: str, optimization_parameters: Dict): return self._post(f"workspaces/{workspace_id}/optimizations", json=optimization_parameters) - def run_optimization_trial( - self, optimization_id: str, geometry_id: str, geometry_parameters: Dict - ): + def run_optimization_trial(self, optimization_id: str, parameters: Dict): return self._post( - f"optimizations/{optimization_id}/trial-runs/{geometry_id}", - json=geometry_parameters, + f"optimizations/{optimization_id}/trial-runs", + json=parameters, ) def get_optimization(self, optimization_id: str): diff --git a/src/ansys/simai/core/data/optimizations.py b/src/ansys/simai/core/data/optimizations.py index 179cd291..bef2874a 100644 --- a/src/ansys/simai/core/data/optimizations.py +++ b/src/ansys/simai/core/data/optimizations.py @@ -22,7 +22,7 @@ import logging from inspect import signature -from typing import Callable, Dict, List, Optional, Tuple +from typing import Callable, Dict, List, Literal, Optional, Tuple from tqdm import tqdm from wakepy import keep @@ -33,6 +33,7 @@ Identifiable, NamedFile, get_id_from_identifiable, + get_object_from_identifiable, ) from ansys.simai.core.data.workspaces import Workspace from ansys.simai.core.errors import InvalidArguments @@ -43,12 +44,8 @@ class Optimization(ComputableDataModel): """Provides the local representation of an optimization definition object.""" - def _try_geometry( - self, geometry: Identifiable[Geometry], geometry_parameters: Dict - ) -> "OptimizationTrialRun": - return self._client._optimization_trial_run_directory._try_geometry( - self.id, geometry, geometry_parameters - ) + def _run_iteration(self, parameters: Dict) -> "OptimizationTrialRun": + return self._client._optimization_trial_run_directory._run_iteration(self.id, parameters) class OptimizationTrialRun(ComputableDataModel): @@ -86,7 +83,7 @@ def get(self, optimization_id: str) -> Optimization: """ return self._model_from(self._client._api.get_optimization(optimization_id)) - def run( + def run_parametric( self, geometry_generation_fn: Callable[..., NamedFile], geometry_parameters: Dict[str, Tuple[float, float]], @@ -98,7 +95,7 @@ def run( show_progress: bool = False, workspace: Optional[Identifiable[Workspace]] = None, ) -> List[Dict]: - """Run an optimization process. + """Run an optimization loop with client-side parametric geometry generation. Args: geometry_generation_fn: Function to call to generate a new geometry @@ -107,6 +104,7 @@ def run( geometry_parameters: Name of the geometry parameters and their bounds or possible values (choices). boundary_conditions: Values of the boundary conditions to perform the optimization at. The values should map to existing boundary conditions in your project/workspace. + n_iters: Number of iterations of the optimization loop. minimize: List of global coefficients to minimize. The global coefficients should map to existing coefficients in your project/workspace. maximize: List of global coefficients to maximize. @@ -119,7 +117,6 @@ def run( - ``x`` is a float bound. - The comparison operator is ``>=`` or ``<=``. - n_iters: Number of iterations of the optimization loop. show_progress: Whether to print progress on stdout. workspace: Workspace to run the optimization in. If a workspace is not specified, the default is the configured workspace. @@ -162,22 +159,18 @@ def my_geometry_generation_function(param_a, param_b): print(results) """ workspace_id = get_id_from_identifiable(workspace, True, self._client._current_workspace) - _check_geometry_generation_fn_signature(geometry_generation_fn, geometry_parameters) - if not minimize and not maximize: - raise InvalidArguments("No global coefficient to optimize.") - objective = {} - if minimize: - for global_coefficient in minimize: - objective[global_coefficient] = {"minimize": True} - if maximize: - for global_coefficient in maximize: - objective[global_coefficient] = {"minimize": False} + _validate_geometry_parameters(geometry_parameters) + _validate_geometry_generation_fn_signature(geometry_generation_fn, geometry_parameters) + objective = _build_objective(minimize, maximize) optimization_parameters = { "boundary_conditions": boundary_conditions, - "geometry_parameters": geometry_parameters, "n_iters": n_iters, "objective": objective, + "type": "parametric", "outcome_constraints": outcome_constraints or [], + "geometry_generation": { + "geometry_parameters": geometry_parameters, + }, } with tqdm(total=n_iters, disable=not show_progress) as progress_bar: progress_bar.set_description("Creating optimization definition.") @@ -204,7 +197,9 @@ def my_geometry_generation_function(param_a, param_b): ) logger.debug("Running trial.") progress_bar.set_description("Running trial.") - trial_run = optimization._try_geometry(geometry, geometry_parameters) + trial_run = optimization._run_iteration( + {"geometry": geometry.id, "geometry_parameters": geometry_parameters} + ) trial_run.wait() iteration_result = { "parameters": geometry_parameters, @@ -222,8 +217,102 @@ def my_geometry_generation_function(param_a, param_b): progress_bar.set_description("Optimization complete.") return iterations_results + def run_non_parametric( + self, + geometry: Identifiable[Geometry], + bounding_boxes: List[List[float]], + symmetries: List[Literal["x", "y", "z", "X", "Y", "Z"]], + boundary_conditions: Dict[str, float], + n_iters: int, + minimize: Optional[List[str]] = None, + maximize: Optional[List[str]] = None, + show_progress: bool = False, + ): + """Run an optimization loop with server-side geometry generation using automorphism. + + Args: + geometry: The base geometry on which to perform the automorphism. The optimization will + run in the same workspace as the geometry. + bounding_boxes: list of the bounds of the different boxes that will define the locations + of the geometry to optimize. + The format is [ + [box1_xmin, box1_xmax, box1_ymin, box1_ymax, box1_zmin, box1_zmax], + [box2_xmin, box2_xmax, box2_ymin, box2_ymax, box2_zmin, box2_zmax], + ... + ] + symmetries: list of symmetry axes, axes being x, y or z + boundary_conditions: Values of the boundary conditions to perform the optimization at. + The values should map to existing boundary conditions in your project/workspace. + n_iters: Number of iterations of the optimization loop. + minimize: List of global coefficients to minimize. + The global coefficients should map to existing coefficients in your project/workspace. + maximize: List of global coefficients to maximize. + The global coefficients should map to existing coefficients in your project/workspace. + show_progress: Whether to print progress on stdout. + workspace: Workspace to run the optimization in. If a workspace is + not specified, the default is the configured workspace. + """ + geometry = get_object_from_identifiable(geometry, self._client._geometry_directory) + objective = _build_objective(minimize, maximize) + optimization_parameters = { + "boundary_conditions": boundary_conditions, + "n_iters": n_iters, + "objective": objective, + "type": "non_parametric", + "geometry_generation": { + "geometry": geometry.id, + "box_bounds_list": bounding_boxes, + "symmetries": symmetries, + }, + } + with tqdm(total=n_iters, disable=not show_progress) as progress_bar: + progress_bar.set_description("Creating optimization definition.") + optimization = self._model_from( + self._client._api.define_optimization( + geometry._fields["workspace_id"], optimization_parameters + ) + ) + optimization.wait() + logger.debug("Optimization defined. Starting optimization loop.") + iterations_results: List[Dict] = [] + with keep.running(on_fail="warn"): + for _ in range(n_iters): + logger.debug("Running iteration") + progress_bar.set_description("Running iteration") + trial_run = optimization._run_iteration({}) + trial_run.wait() + iteration_result = { + "objective": trial_run.fields["outcome_values"], + } + progress_bar.set_postfix(**iteration_result) + if trial_run.fields.get("is_feasible", True): + iterations_results.append(iteration_result) + else: + logger.debug("Trial run results did not match constraints. Skipping.") + logger.debug("Trial complete.") + progress_bar.update(1) + logger.debug("Optimization complete.") + progress_bar.set_description("Optimization complete.") + return iterations_results + -def _check_geometry_generation_fn_signature(geometry_generation_fn, geometry_parameters): +def _validate_geometry_parameters(params: Dict): + if not isinstance(params, Dict): + raise InvalidArguments("geometry_parameters: must be a dict.") + if not params: + raise InvalidArguments("geometry_parameters: must not be empty.") + for key, value in params.items(): + bounds = value.get("bounds") + choices = value.get("choices") + if not bounds and not choices: + raise InvalidArguments(f"geometry_parameters: no bounds or choices specified for {key}") + if bounds and choices: + raise InvalidArguments( + f"geometry_parameters: only one of bounds or choices must be specified for {key}" + ) + + +def _validate_geometry_generation_fn_signature(geometry_generation_fn, geometry_parameters): geometry_generation_fn_args = signature(geometry_generation_fn).parameters if geometry_generation_fn_args.keys() != geometry_parameters.keys(): raise InvalidArguments( @@ -231,6 +320,19 @@ def _check_geometry_generation_fn_signature(geometry_generation_fn, geometry_par ) +def _build_objective(minimize: List[str], maximize: List[str]) -> Dict: + if not minimize and not maximize: + raise InvalidArguments("No global coefficient to optimize.") + objective = {} + if minimize: + for global_coefficient in minimize: + objective[global_coefficient] = {"minimize": True} + if maximize: + for global_coefficient in maximize: + objective[global_coefficient] = {"minimize": False} + return objective + + # Undocumented for now, users don't really need to interact with it class OptimizationTrialRunDirectory(Directory[OptimizationTrialRun]): _data_model = OptimizationTrialRun @@ -239,16 +341,10 @@ def get(self, trial_run_id: str): """Get a specific trial run from the server.""" return self._model_from(self._client._api.get_optimization_trial_run(trial_run_id)) - def _try_geometry( - self, - optimization: Identifiable[Optimization], - geometry: Identifiable[Geometry], - geometry_parameters: Dict, + def _run_iteration( + self, optimization: Identifiable[Optimization], parameters: Dict ) -> OptimizationTrialRun: - geometry_id = get_id_from_identifiable(geometry) optimization_id = get_id_from_identifiable(optimization) return self._model_from( - self._client._api.run_optimization_trial( - optimization_id, geometry_id, geometry_parameters - ) + self._client._api.run_optimization_trial(optimization_id, parameters) ) diff --git a/tests/test_optimizations.py b/tests/test_optimizations.py index 6762c297..bf7ab06a 100644 --- a/tests/test_optimizations.py +++ b/tests/test_optimizations.py @@ -21,9 +21,14 @@ # SOFTWARE. # ruff: noqa: E731 +import json +import threading + import pytest +import responses +import sseclient -from ansys.simai.core.data.optimizations import _check_geometry_generation_fn_signature +from ansys.simai.core.data.optimizations import _validate_geometry_generation_fn_signature from ansys.simai.core.errors import InvalidArguments @@ -39,7 +44,7 @@ def test_geometry_generation_fn_invalid_signature(simai_client): } with pytest.raises(InvalidArguments) as exc: - simai_client.optimizations.run( + simai_client.optimizations.run_parametric( geometry_generation_fn=my_geometry_generation_function, geometry_parameters=geometry_parameters, boundary_conditions={"abc": 3.0}, @@ -49,7 +54,7 @@ def test_geometry_generation_fn_invalid_signature(simai_client): assert "geometry_generation_fn requires the following signature" in str(exc.value) -def test_check_geometry_generation_fn_valid_signature(): +def test_validate_geometry_generation_fn_valid_signature(): """WHEN geometry_generation_fn signature matches geometry_parameters keys THEN check passes""" @@ -58,4 +63,176 @@ def test_check_geometry_generation_fn_valid_signature(): ) geometry_parameters = {"param_c": {"bounds": (-12.5, 12.5)}, "param_d": {"choices": (0.1, 1.0)}} - _check_geometry_generation_fn_signature(my_geometry_generation_function, geometry_parameters) + _validate_geometry_generation_fn_signature(my_geometry_generation_function, geometry_parameters) + + +@responses.activate +def test_run_parametric_optimization(simai_client, mocker): + workspace_id = "insert_cool_reference" + mocker.patch("ansys.simai.core.data.optimizations._validate_geometry_generation_fn_signature") + geometry_upload = mocker.patch.object(simai_client.geometries, "upload") + geometry_upload.return_value = simai_client.geometries._model_from({"id": "geomx"}) + + responses.add( + responses.POST, + f"https://test.test/workspaces/{workspace_id}/optimizations", + status=202, + json={"id": "wow", "state": "requested", "trial_runs": []}, + ) + responses.add( + responses.POST, + "https://test.test/optimizations/wow/trial-runs", + status=202, + json={"id": "wow1", "state": "requested"}, + ) + responses.add( + responses.POST, + "https://test.test/optimizations/wow/trial-runs", + status=202, + json={"id": "wow2", "state": "requested"}, + ) + responses.add( + responses.POST, + "https://test.test/optimizations/wow/trial-runs", + status=202, + json={"id": "wow3", "state": "requested"}, + ) + threading.Timer( + 0.1, + simai_client._api._handle_sse_event, + args=[ + sseclient.Event( + data=json.dumps( + { + "type": "job", + "status": "successful", + "record": { + "id": "wow", + "state": "successful", + "initial_geometry_parameters": {"a": 1}, + }, + "target": {"type": "optimization", "id": "wow"}, + } + ) + ) + ], + ).start() + for i in range(1, 4): + threading.Timer( + i / 10 + 0.1, + simai_client._api._handle_sse_event, + args=[ + sseclient.Event( + data=json.dumps( + { + "type": "job", + "status": "successful", + "record": { + "id": f"wow{i}", + "state": "successful", + "is_feasible": "true", + "outcome_values": i, + "next_geometry_parameters": {"a": i + 1} if i != 3 else None, + }, + "target": {"type": "optimization_trial_run", "id": f"wow{i}"}, + } + ) + ) + ], + ).start() + geometry_generation_fn = mocker.stub(name="geometry_gen") + results = simai_client.optimizations.run_parametric( + geometry_generation_fn=geometry_generation_fn, + geometry_parameters={ + "a": {"bounds": (-12.5, 12.5)}, + }, + minimize=["TotalForceX"], + boundary_conditions={"VelocityX": 10.5}, + outcome_constraints=["TotalForceX <= 10"], + n_iters=3, + workspace=workspace_id, + ) + assert results == [{"objective": i, "parameters": {"a": i}} for i in range(1, 4)] + + +@responses.activate +def test_run_non_parametric_optimization(simai_client, geometry_factory): + workspace_id = "insert_cool_reference" + geometry = geometry_factory(workspace_id=workspace_id) + + responses.add( + responses.POST, + f"https://test.test/workspaces/{workspace_id}/optimizations", + status=202, + json={"id": "wow", "state": "requested", "trial_runs": []}, + ) + responses.add( + responses.POST, + "https://test.test/optimizations/wow/trial-runs", + status=202, + json={"id": "wow1", "state": "requested"}, + ) + responses.add( + responses.POST, + "https://test.test/optimizations/wow/trial-runs", + status=202, + json={"id": "wow2", "state": "requested"}, + ) + responses.add( + responses.POST, + "https://test.test/optimizations/wow/trial-runs", + status=202, + json={"id": "wow3", "state": "requested"}, + ) + threading.Timer( + 0.1, + simai_client._api._handle_sse_event, + args=[ + sseclient.Event( + data=json.dumps( + { + "type": "job", + "status": "successful", + "record": { + "id": "wow", + "state": "successful", + "initial_geometry_parameters": None, + }, + "target": {"type": "optimization", "id": "wow"}, + } + ) + ) + ], + ).start() + for i in range(1, 4): + threading.Timer( + i / 10 + 0.1, + simai_client._api._handle_sse_event, + args=[ + sseclient.Event( + data=json.dumps( + { + "type": "job", + "status": "successful", + "record": { + "id": f"wow{i}", + "state": "successful", + "is_feasible": "true", + "outcome_values": i, + "next_geometry_parameters": None, + }, + "target": {"type": "optimization_trial_run", "id": f"wow{i}"}, + } + ) + ) + ], + ).start() + results = simai_client.optimizations.run_non_parametric( + geometry=geometry, + bounding_boxes=[[1, 1, 1, 1, 1, 1]], + symmetries=["x", "y", "z"], + minimize=["TotalForceX"], + boundary_conditions={"VelocityX": 10.5}, + n_iters=3, + ) + assert results == [{"objective": i} for i in range(1, 4)]