Enabled all output types to use UserVariable

flow360/component/simulation/outputs/outputs.py

@@ -194,7 +194,7 @@ class SurfaceOutput(_AnimationAndFileFormatSettings):
         + "Will choose the value of the last instance of this option of the same output type "
         + "(:class:`SurfaceOutput` or :class:`TimeAverageSurfaceOutput`) in the output list.",
     )
-    output_fields: UniqueItemList[Union[SurfaceFieldNames, str]] = pd.Field(
+    output_fields: UniqueItemList[Union[SurfaceFieldNames, str, UserVariable]] = pd.Field(
         description="List of output variables. Including :ref:`universal output variables<UniversalVariablesV2>`,"
         + " :ref:`variables specific to SurfaceOutput<SurfaceSpecificVariablesV2>` and :class:`UserDefinedField`."
     )
@@ -329,7 +329,7 @@ class SliceOutput(_AnimationAndFileFormatSettings):
         alias="slices",
         description="List of output :class:`~flow360.Slice` entities.",
     )
-    output_fields: UniqueItemList[Union[SliceFieldNames, str]] = pd.Field(
+    output_fields: UniqueItemList[Union[SliceFieldNames, str, UserVariable]] = pd.Field(
         description="List of output variables. Including :ref:`universal output variables<UniversalVariablesV2>`,"
         " :ref:`variables specific to SliceOutput<VolumeAndSliceSpecificVariablesV2>`"
         " and :class:`UserDefinedField`."
@@ -415,7 +415,7 @@ class IsosurfaceOutput(_AnimationAndFileFormatSettings):
         alias="isosurfaces",
         description="List of :class:`~flow360.Isosurface` entities.",
     )
-    output_fields: UniqueItemList[Union[CommonFieldNames, str]] = pd.Field(
+    output_fields: UniqueItemList[Union[CommonFieldNames, str, UserVariable]] = pd.Field(
         description="List of output variables. Including "
         ":ref:`universal output variables<UniversalVariablesV2>` and :class:`UserDefinedField`."
     )
@@ -452,7 +452,7 @@ class SurfaceIntegralOutput(_OutputBase):
         alias="surfaces",
         description="List of boundaries where the surface integral will be calculated.",
     )
-    output_fields: UniqueItemList[str] = pd.Field(
+    output_fields: UniqueItemList[Union[str, UserVariable]] = pd.Field(
         description="List of output variables, only the :class:`UserDefinedField` is allowed."
     )
     output_type: Literal["SurfaceIntegralOutput"] = pd.Field("SurfaceIntegralOutput", frozen=True)
@@ -516,7 +516,7 @@ class ProbeOutput(_OutputBase):
         + "monitor group. :class:`~flow360.PointArray` is used to "
         + "define monitored points along a line.",
     )
-    output_fields: UniqueItemList[Union[CommonFieldNames, str]] = pd.Field(
+    output_fields: UniqueItemList[Union[CommonFieldNames, str, UserVariable]] = pd.Field(
         description="List of output fields. Including :ref:`universal output variables<UniversalVariablesV2>`"
         " and :class:`UserDefinedField`."
     )
@@ -580,7 +580,7 @@ class SurfaceProbeOutput(Flow360BaseModel):
         + "entities belonging to this monitor group."
     )
 
-    output_fields: UniqueItemList[Union[SurfaceFieldNames, str]] = pd.Field(
+    output_fields: UniqueItemList[Union[SurfaceFieldNames, str, UserVariable]] = pd.Field(
         description="List of output variables. Including :ref:`universal output variables<UniversalVariablesV2>`,"
         " :ref:`variables specific to SurfaceOutput<SurfaceSpecificVariablesV2>` and :class:`UserDefinedField`."
     )
@@ -609,7 +609,7 @@ class SurfaceSliceOutput(_AnimationAndFileFormatSettings):
 
     output_format: Literal["paraview"] = pd.Field(default="paraview")
 
-    output_fields: UniqueItemList[Union[SurfaceFieldNames, str]] = pd.Field(
+    output_fields: UniqueItemList[Union[SurfaceFieldNames, str, UserVariable]] = pd.Field(
         description="List of output variables. Including :ref:`universal output variables<UniversalVariablesV2>`,"
         " :ref:`variables specific to SurfaceOutput<SurfaceSpecificVariablesV2>` and :class:`UserDefinedField`."
     )

flow360/component/simulation/translator/solver_translator.py

@@ -249,8 +249,10 @@ def translate_output_fields(
 
     for output_field in output_model.output_fields.items:
         if isinstance(output_field, UserVariable):
+            # Remove the UserVariable object and add its name
             output_fields.append(output_field.name)
-
+    # Filter out the UserVariable Dicts
+    output_fields = [item for item in output_fields if isinstance(item, str)]
     return {"outputFields": output_fields}
 
 
@@ -382,6 +384,8 @@ def translate_volume_output(
     ).items:
         if isinstance(output_field, UserVariable):
             output_fields.append(output_field.name)
+    # Filter out the UserVariable Dicts
+    output_fields = [item for item in output_fields if isinstance(item, str)]
     volume_output.update(
         {
             "outputFields": output_fields,
@@ -673,18 +677,18 @@ def process_output_fields_for_udf(input_params: SimulationParams):
         if udf_expression:
             generated_udfs.append(UserDefinedField(name=field_name, expression=udf_expression))
 
+    # UserVariable handling:
+    user_variable_udfs = {}
     if input_params.outputs:
-        # UserVariable handling:
-        user_variable_udfs = set()
         for output in input_params.outputs:
             if not hasattr(output, "output_fields") or not output.output_fields:
                 continue
             for output_field in output.output_fields.items:
                 if not isinstance(output_field, UserVariable):
                     continue
-                user_variable_udfs.add(user_variable_to_udf(output_field, input_params))
-
-    return generated_udfs
+                udf_from_user_variable = user_variable_to_udf(output_field, input_params)
+                user_variable_udfs[udf_from_user_variable.name] = udf_from_user_variable
+    return generated_udfs + list(user_variable_udfs.values())
 
 
 def translate_streamline_output(output_params: list):

flow360/component/simulation/user_code/core/types.py

@@ -308,14 +308,15 @@ def update_context(self):
             default_context.set_alias(self.name, self.solver_name)
         return self
 
-    def in_unit(self, new_name: str, new_unit: str = None):
+    def in_unit(self, new_name: str, new_unit: Union[str, Unit] = None):
         """
 
 
         Return a UserVariable that will generate results in the new_unit.
         If new_unit is not specified then the unit will be determined by the unit system.
         """
-
+        if isinstance(new_unit, Unit):
+            new_unit = str(new_unit)
         new_variable = UserVariable(name=new_name, value=Expression(expression=self.name))
         new_variable.value.output_units = new_unit  # pylint:disable=assigning-non-slot
         return new_variable