[FL-727][FLPY-3] Support group by boundary and by body in surface force csv result model

flow360/component/results/case_results.py

@@ -1,4 +1,4 @@
-""" Case results module"""
+"""Case results module"""
 
 from __future__ import annotations
 
@@ -8,6 +8,7 @@
 import tempfile
 import time
 import uuid
+from collections import defaultdict
 from enum import Enum
 from itertools import chain, product
 from typing import Callable, Dict, List, Optional
@@ -31,6 +32,9 @@
 from ...exceptions import Flow360ValueError
 from ...log import log
 from ..simulation.conversion import unit_converter as unit_converter_v2
+from ..simulation.entity_info import GeometryEntityInfo
+from ..simulation.models.surface_models import BoundaryBase
+from ..simulation.primitives import GeometryBodyGroup, Surface
 from ..simulation.simulation_params import SimulationParams
 from ..v1.conversions import unit_converter as unit_converter_v1
 from ..v1.flow360_params import Flow360Params
@@ -226,6 +230,74 @@ def _preprocess(self, filter_physical_steps_only: bool = True, include_time: boo
     def reload_data(self, filter_physical_steps_only: bool = True, include_time: bool = True):
         return super().reload_data(filter_physical_steps_only, include_time)
 
+    def _create_surface_forces_group(
+        self, entity_groups: Dict[str, List[str]]
+    ) -> SurfaceForcesGroupResultCSVModel:
+        """
+        Create the SurfaceForcesGroupResultCSVModel for the given entity groups.
+        """
+        raw_values = {}
+        for x_column in self._x_columns:
+            raw_values[x_column] = np.array(self.raw_values[x_column])
+        for name, entities in entity_groups.items():
+            self.filter(include=entities)
+            for variable in self._variables:
+                if f"total{variable}" not in raw_values:
+                    raw_values[f"{name}_{variable}"] = np.array(self.values[f"total{variable}"])
+                    continue
+                raw_values[f"{name}_{variable}"] += np.array(self.values[f"total{variable}"])
+
+        raw_values = {key: val.tolist() for key, val in raw_values.items()}
+        entity_groups = {key: sorted(val) for key, val in entity_groups.items()}
+
+        return SurfaceForcesGroupResultCSVModel.from_dict(data=raw_values, group=entity_groups)
+
+    def by_boundary_condition(self, params: SimulationParams) -> SurfaceForcesGroupResultCSVModel:
+        """
+        Group entities by boundary condition's name and create a
+        SurfaceForcesGroupResultCSVModel
+        """
+
+        entity_groups = defaultdict(list)
+        for model in params.models:
+            if not isinstance(model, BoundaryBase):
+                continue
+            boundary_name = model.name
+            if boundary_name is None:
+                boundary_name = model.type
+            entity_groups[boundary_name].extend(
+                [entity.name for entity in model.entities.stored_entities]
+            )
+        return self._create_surface_forces_group(entity_groups=entity_groups)
+
+    def by_body_group(self, params: SimulationParams) -> SurfaceForcesGroupResultCSVModel:
+        """
+        Group entities by body group's name and create a
+        SurfaceForcesGroupResultCSVModel
+        """
+        if not isinstance(
+            params.private_attribute_asset_cache.project_entity_info, GeometryEntityInfo
+        ):
+            raise Flow360ValueError(
+                "Group surface forces by body group is only supported for case starting from geometry."
+            )
+        entity_info = params.private_attribute_asset_cache.project_entity_info
+        entity_groups = entity_info.get_body_group_to_face_group_name_map()
+        return self._create_surface_forces_group(entity_groups=entity_groups)
+
+
+class SurfaceForcesGroupResultCSVModel(SurfaceForcesResultCSVModel):
+    """SurfaceForcesGroupResultCSVModel"""
+
+    remote_file_name: str = pd.Field(None, frozen=True)
+    entity_groups: Optional[dict] = pd.Field(None)
+
+    @classmethod
+    def from_dict(cls, data: dict, group: dict = None):
+        obj = super().from_dict(data)
+        obj.entity_groups = group
+        return obj
+
 
 class LegacyForceDistributionResultCSVModel(ResultCSVModel):
     """ForceDistributionResultCSVModel"""

flow360/component/simulation/entity_info.py

@@ -1,6 +1,7 @@
 """Deserializer for entity info retrieved from asset metadata pipeline."""
 
 from abc import ABCMeta, abstractmethod
+from collections import defaultdict
 from typing import Annotated, List, Literal, Optional, Union
 
 import pydantic as pd
@@ -411,6 +412,65 @@ def compute_transformation_matrices(self):
                 body_group.transformation.get_transformation_matrix().flatten().tolist()
             )
 
+    def get_body_group_to_face_group_name_map(self) -> dict[str, list[str]]:
+        """
+        Returns bodyId to file name mapping.
+        """
+
+        # pylint: disable=too-many-locals
+        def create_group_to_sub_component_mapping(group):
+            mapping = defaultdict(list)
+            for item in group:
+                mapping[item.private_attribute_id].extend(item.private_attribute_sub_components)
+            return mapping
+
+        body_group_name_to_id_current = create_group_to_sub_component_mapping(
+            self._get_list_of_entities(entity_type_name="body", attribute_name=self.body_group_tag)
+        )
+        face_group_name_to_id_current = create_group_to_sub_component_mapping(
+            self._get_list_of_entities(entity_type_name="face", attribute_name=self.face_group_tag)
+        )
+        face_group_name_to_id_by_body_id = create_group_to_sub_component_mapping(
+            self._get_list_of_entities(entity_type_name="face", attribute_name="groupByBodyId")
+        )
+
+        body_group_name_to_face_id = defaultdict(list)
+        for body_group_name, body_ids in body_group_name_to_id_current.items():
+            for body_id in body_ids:
+                body_group_name_to_face_id[body_group_name].extend(
+                    face_group_name_to_id_by_body_id[body_id]
+                )
+
+        face_id_to_body_group_name = {}
+        for body_group_name, face_ids in body_group_name_to_face_id.items():
+            for face_id in face_ids:
+                face_id_to_body_group_name[face_id] = body_group_name
+
+        body_group_name_to_face_group_name = defaultdict(list)
+        for fact_group_name, face_ids in face_group_name_to_id_current.items():
+            body_group_in_this_face_group = set()
+            for face_id in face_ids:
+                owning_body = face_id_to_body_group_name.get(face_id)
+                if owning_body is None:
+                    # A face in a face group was not found in any body group
+                    raise ValueError(
+                        f"Face ID '{face_id}' found in face group '{fact_group_name}' "
+                        "but not found in any body group."
+                    )
+                body_group_in_this_face_group.add(owning_body)
+            if len(body_group_in_this_face_group) > 1:
+                raise ValueError(
+                    f"Face group '{fact_group_name}' contains faces belonging to multiple body groups: "
+                    f"{list(sorted(body_group_in_this_face_group))}. "
+                    "The mapping between body and face groups cannot be created."
+                )
+
+            if len(body_group_in_this_face_group) == 1:
+                owning_body = list(body_group_in_this_face_group)[0]
+                body_group_name_to_face_group_name[owning_body].append(fact_group_name)
+
+        return body_group_name_to_face_group_name
+
 
 class VolumeMeshEntityInfo(EntityInfoModel):
     """Data model for volume mesh entityInfo.json"""