Search space: Improving constraint handling

CHANGELOG.md

@@ -6,11 +6,15 @@
 - Parameters in ParamSpace can also be indexed by name
 - Parameters now have search_space property, to modify the optimizer search space from the full space
 - Continuous parameters have search_min/search_max; Discete parameteres have search_categories
+- Constraints are now defined by Constraint class
+- Input constraints can now be included in ParamSpace, and serialized from there
+- Output constraints can now be included in Campaign, and serialized from there
+- New interface class IParamSpace to address circular import issues between ParamSpace and Constraint
 
 ### Modified
 - Optimizer and Campaign X_space attributes are now assigned using setter
 
-### Remvoed
+### Removed
 - Torch device references and options (GPU compatibility may be re-added)
 
 ## [0.8.4]

obsidian/campaign/campaign.py

@@ -4,7 +4,9 @@
 from obsidian.optimizer import Optimizer, BayesianOptimizer
 from obsidian.experiment import ExpDesigner
 from obsidian.objectives import Objective, Objective_Sequence, obj_class_dict
+from obsidian.constraints import Output_Constraint, const_class_dict
 from obsidian.exceptions import IncompatibleObjectiveError
+from obsidian.utils import tensordict_to_dict
 import obsidian
 
 import pandas as pd
@@ -42,6 +44,7 @@ class Campaign():
     def __init__(self,
                  X_space: ParamSpace,
                  target: Target | list[Target],
+                 constraints: Output_Constraint | list[Output_Constraint] | None = None,
                  optimizer: Optimizer | None = None,
                  designer: ExpDesigner | None = None,
                  objective: Objective | None = None,
@@ -59,6 +62,9 @@ def __init__(self,
         self.set_target(target)
         self.set_objective(objective)
 
+        self.output_constraints = None
+        self.constrain_outputs(constraints)
+
         # Non-object attributes
         self.iter = 0
         self.seed = seed
@@ -269,59 +275,6 @@ def X(self) -> pd.DataFrame:
         """
         return self.data[list(self.X_space.X_names)]
 
-    def save_state(self) -> dict:
-        """
-        Saves the state of the Campaign object as a dictionary.
-
-        Returns:
-            dict: A dictionary containing the saved state of the Campaign object.
-        """
-        obj_dict = {}
-        obj_dict['X_space'] = self.X_space.save_state()
-        obj_dict['optimizer'] = self.optimizer.save_state()
-        obj_dict['data'] = self.data.to_dict()
-        obj_dict['target'] = [t.save_state() for t in self.target]
-        if self.objective:
-            obj_dict['objective'] = self.objective.save_state()
-        obj_dict['seed'] = self.seed
-
-        return obj_dict
-
-    @classmethod
-    def load_state(cls,
-                   obj_dict: dict):
-        """
-        Loads the state of the campaign from a dictionary.
-
-        Args:
-            cls (Campaign): The class object.
-            obj_dict (dict): A dictionary containing the campaign state.
-
-        Returns:
-            Campaign: A new campaign object with the loaded state.
-        """
-
-        if 'objective' in obj_dict:
-            if obj_dict['objective']['name'] == 'Objective_Sequence':
-                new_objective = Objective_Sequence.load_state(obj_dict['objective'])
-            else:
-                obj_class = obj_class_dict[obj_dict['objective']['name']]
-                new_objective = obj_class.load_state(obj_dict['objective'])
-        else:
-            new_objective = None
-
-        new_campaign = cls(X_space=ParamSpace.load_state(obj_dict['X_space']),
-                           target=[Target.load_state(t_dict) for t_dict in obj_dict['target']],
-                           optimizer=BayesianOptimizer.load_state(obj_dict['optimizer']),
-                           objective=new_objective,
-                           seed=obj_dict['seed'])
-        new_campaign.data = pd.DataFrame(obj_dict['data'])
-        new_campaign.data.index = new_campaign.data.index.astype('int')
-
-        new_campaign.iter = new_campaign.data['Iteration'].astype('int').max()
-
-        return new_campaign
-
     def __repr__(self):
         """String representation of object"""
         return f"obsidian Campaign for {getattr(self,'y_names', None)}; {getattr(self,'m_exp', 0)} observations"
@@ -353,7 +306,9 @@ def suggest(self, **optim_kwargs):
             try:
                 # In case X_space has changed, re-set the optimizer X_space
                 self.optimizer.set_X_space(self.X_space)
-                X, eval = self.optimizer.suggest(objective=self.objective, **optim_kwargs)
+                X, eval = self.optimizer.suggest(objective=self.objective,
+                                                 out_constraints=self.output_constraints,
+                                                 **optim_kwargs)
                 return (X, eval)
             except Exception:
                 warnings.warn('Optimization failed')
@@ -430,3 +385,82 @@ def _analyze(self):
             )
 
         return
+
+    def constrain_outputs(self,
+                          constraints: Output_Constraint | list[Output_Constraint] | None) -> None:
+        """
+        Sets optional output constraints for the campaign.
+        """
+        if constraints is not None:
+            if isinstance(constraints, Output_Constraint):
+                constraints = [constraints]
+            self.output_constraints = constraints
+
+        return
+
+    def clear_output_constraints(self):
+        """Clears output constraints"""
+        self.output_constraints = None
+
+    def save_state(self) -> dict:
+        """
+        Saves the state of the Campaign object as a dictionary.
+
+        Returns:
+            dict: A dictionary containing the saved state of the Campaign object.
+        """
+        obj_dict = {}
+        obj_dict['X_space'] = self.X_space.save_state()
+        obj_dict['optimizer'] = self.optimizer.save_state()
+        obj_dict['data'] = self.data.to_dict()
+        obj_dict['target'] = [t.save_state() for t in self.target]
+        if self.objective:
+            obj_dict['objective'] = self.objective.save_state()
+        obj_dict['seed'] = self.seed
+
+        if getattr(self, 'output_constraints', None):
+            obj_dict['output_constraints'] = [{'class': const.__class__.__name__,
+                                               'state': tensordict_to_dict(const.state_dict())}
+                                              for const in self.output_constraints]
+
+        return obj_dict
+
+    @classmethod
+    def load_state(cls,
+                   obj_dict: dict):
+        """
+        Loads the state of the campaign from a dictionary.
+
+        Args:
+            cls (Campaign): The class object.
+            obj_dict (dict): A dictionary containing the campaign state.
+
+        Returns:
+            Campaign: A new campaign object with the loaded state.
+        """
+
+        if 'objective' in obj_dict:
+            if obj_dict['objective']['name'] == 'Objective_Sequence':
+                new_objective = Objective_Sequence.load_state(obj_dict['objective'])
+            else:
+                obj_class = obj_class_dict[obj_dict['objective']['name']]
+                new_objective = obj_class.load_state(obj_dict['objective'])
+        else:
+            new_objective = None
+
+        new_campaign = cls(X_space=ParamSpace.load_state(obj_dict['X_space']),
+                           target=[Target.load_state(t_dict) for t_dict in obj_dict['target']],
+                           optimizer=BayesianOptimizer.load_state(obj_dict['optimizer']),
+                           objective=new_objective,
+                           seed=obj_dict['seed'])
+        new_campaign.data = pd.DataFrame(obj_dict['data'])
+        new_campaign.data.index = new_campaign.data.index.astype('int')
+
+        new_campaign.iter = new_campaign.data['Iteration'].astype('int').max()
+
+        if 'output_constraints' in obj_dict:
+            for const_dict in obj_dict['output_constraints']:
+                const = const_class_dict[const_dict['class']](new_campaign.target, **const_dict['state'])
+                new_campaign.constrain_outputs(const)
+
+        return new_campaign

obsidian/constraints/__init__.py

@@ -1,4 +1,6 @@
 """Constraints: Restrict the recommended space during optimization"""
 
+from .base import *
 from .input import *
 from .output import *
+from .config import *

obsidian/constraints/base.py

@@ -0,0 +1,20 @@
+"""Base class for obsidian constraints"""
+
+
+from abc import abstractmethod, ABC
+from torch.nn import Module
+
+
+class Constraint(ABC, Module):
+    """
+    Base class for constraints, which restrict the input or output space
+    of a model or optimization problem
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+        return
+
+    @abstractmethod
+    def forward(self):
+        pass  # pragma: no cover

obsidian/constraints/config.py

@@ -0,0 +1,17 @@
+"""Method pointers and config for constraints"""
+
+from .input import (
+    Linear_Constraint,
+    BatchVariance_Constraint
+)
+
+from .output import (
+    Blank_Constraint,
+    L1_Constraint
+)
+
+const_class_dict = {'Linear_Constraint': Linear_Constraint,
+                    'BatchVariance_Constraint': BatchVariance_Constraint,
+                    'Blank_Constraint': Blank_Constraint,
+                    'L1_Constraint': L1_Constraint
+                    }