Add scalers and testing for thickener & dewaterer

watertap/unit_models/dewatering.py

@@ -34,16 +34,18 @@
 import idaes.logger as idaeslog
 
 from pyomo.environ import (
+    Constraint,
     Param,
-    units as pyunits,
     Var,
     NonNegativeReals,
+    units as pyunits,
 )
 from pyomo.common.config import ConfigValue, In
 
 from idaes.core.util.exceptions import (
     ConfigurationError,
 )
+from idaes.core.scaling import CustomScalerBase, ConstraintScalingScheme
 from watertap.costing.unit_models.dewatering import cost_dewatering
 
 __author__ = "Alejandro Garciadiego, Adam Atia"
@@ -65,12 +67,116 @@ class ActivatedSludgeModelType(Enum):
     modified_ASM2D = auto()
 
 
+class DewatererScaler(CustomScalerBase):
+    """
+    Default modular scaler for the dewatering unit model.
+    This Scaler relies on the associated property and reaction packages,
+    either through user provided options (submodel_scalers argument) or by default
+    Scalers assigned to the packages.
+    """
+
+    DEFAULT_SCALING_FACTORS = {
+        "volume": 1e-3,
+    }
+
+    def variable_scaling_routine(
+        self, model, overwrite: bool = False, submodel_scalers: dict = None
+    ):
+        """
+        Routine to apply scaling factors to variables in model.
+        Args:
+            model: model to be scaled
+            overwrite: whether to overwrite existing scaling factors
+            submodel_scalers: dict of Scalers to use for sub-models, keyed by submodel local name
+        Returns:
+            None
+        """
+        # Call scaling methods for sub-models
+        self.call_submodel_scaler_method(
+            submodel=model.mixed_state,
+            method="variable_scaling_routine",
+            submodel_scalers=submodel_scalers,
+            overwrite=overwrite,
+        )
+        self.propagate_state_scaling(
+            target_state=model.underflow_state,
+            source_state=model.mixed_state,
+            overwrite=overwrite,
+        )
+        self.propagate_state_scaling(
+            target_state=model.overflow_state,
+            source_state=model.mixed_state,
+            overwrite=overwrite,
+        )
+
+        self.call_submodel_scaler_method(
+            submodel=model.underflow_state,
+            method="variable_scaling_routine",
+            submodel_scalers=submodel_scalers,
+            overwrite=overwrite,
+        )
+        self.call_submodel_scaler_method(
+            submodel=model.overflow_state,
+            method="variable_scaling_routine",
+            submodel_scalers=submodel_scalers,
+            overwrite=overwrite,
+        )
+
+        # Scale unit level variables
+        self.scale_variable_by_default(model.volume[0], overwrite=overwrite)
+
+    def constraint_scaling_routine(
+        self, model, overwrite: bool = False, submodel_scalers: dict = None
+    ):
+        """
+        Routine to apply scaling factors to constraints in model.
+        Submodel Scalers are called for the property and reaction blocks. All other constraints
+        are scaled using the inverse maximum scheme.
+        Args:
+            model: model to be scaled
+            overwrite: whether to overwrite existing scaling factors
+            submodel_scalers: dict of Scalers to use for sub-models, keyed by submodel local name
+        Returns:
+            None
+        """
+        # Call scaling methods for sub-models
+        self.call_submodel_scaler_method(
+            submodel=model.mixed_state,
+            method="constraint_scaling_routine",
+            submodel_scalers=submodel_scalers,
+            overwrite=overwrite,
+        )
+        self.call_submodel_scaler_method(
+            submodel=model.underflow_state,
+            method="constraint_scaling_routine",
+            submodel_scalers=submodel_scalers,
+            overwrite=overwrite,
+        )
+        self.call_submodel_scaler_method(
+            submodel=model.overflow_state,
+            method="constraint_scaling_routine",
+            submodel_scalers=submodel_scalers,
+            overwrite=overwrite,
+        )
+
+        # Scale unit level constraints
+        # Constraint(
+        for c in model.component_data_objects(Constraint, descend_into=False):
+            self.scale_constraint_by_nominal_value(
+                c,
+                scheme=ConstraintScalingScheme.inverseMaximum,
+                overwrite=overwrite,
+            )
+
+
 @declare_process_block_class("DewateringUnit")
 class DewateringData(SeparatorData):
     """
     Dewatering unit block for BSM2
     """
 
+    default_scaler = DewatererScaler
+
     CONFIG = SeparatorData.CONFIG()
     CONFIG.outlet_list = ["underflow", "overflow"]
     CONFIG.split_basis = SplittingType.componentFlow