luminarycloud 0.13.0__py3-none-any.whl → 0.14.0__py3-none-any.whl

luminarycloud/_helpers/cond.py

@@ -2,11 +2,12 @@
 # type: ignore
 
 import json
-from contextlib import contextmanager
-from dataclasses import dataclass
-from itertools import chain
-from collections.abc import Iterable, Mapping
+from dataclasses import dataclass, field
+from itertools import pairwise
+from collections import defaultdict
+from collections.abc import Iterable, Mapping, Iterator
 from typing import Any, Optional, TypeVar
+from uuid import UUID, uuid4
 
 from google.protobuf.any_pb2 import Any as AnyProto
 from google.protobuf.descriptor import Descriptor, FieldDescriptor
@@ -49,261 +50,353 @@ K = TypeVar("K")
 V = TypeVar("V")
 
 
-class ContextMap(Mapping[K, V]):
+class ContextFrame(Mapping[K, V]):
     """
-    A utility class which implements a Mapping (i.e. dict) interface as a stack of context frames.
+    A utility class which implements a Mapping (i.e. dict) interface as a... graph(?) of context frames.
 
-    Writes only write to the topmost frame, while reads look at each from frame top to bottom and
-    return the first hit. Frames can be pushed and popped onto/from the stack.
+    Writes only write to this frame, while reads look at this frame and then each of its ancestor
+    frames, returning only the first hit. Returns None if there is no hit.
     """
 
-    _frames: list[dict[type[K], type[V]]]
-    _current_frame: dict[type[K], type[V]]
+    __id: UUID
+    __parent_frames: list["ContextFrame"]
+    __current_frame: dict[type[K], type[V]]
 
     def __init__(self):
-        self._current_frame = dict()
-        self._frames = [self._current_frame]
+        self.__id = uuid4()
+        self.__current_frame = dict()
+        self.__parent_frames = []
 
-    @contextmanager
-    def scope(self):
+    def new_frame(self) -> "ContextFrame":
+        """Creates a new frame on top of this frame"""
+        __next_frame = ContextFrame()
+        __next_frame.__parent_frames = [self]
+        return __next_frame
+
+    def entangle(self, other: "ContextFrame") -> None:
         """
-        Pushes a new frame on enter, and pops on exit.
-
-        Examples
-        --------
-        >>> with context_map.scope():
-        >>>     context_map[k] = v
-        >>>     assert context_map[k] == v
-        >>> assert context_map[k] != v
+        Entangles this frame with another frame.
+
+        Essentially, becomes a secondary parent of the other frame, and vice versa. This creates a
+        cycle but reads won't revisit frames. Might seem weird, but it's necessary.
         """
-        self.push()
-        yield self
-        self.pop()
-
-    def push(self) -> None:
-        """Pushes a new frame onto the stack"""
-        self._current_frame = dict()
-        self._frames.append(self._current_frame)
-
-    def pop(self) -> None:
-        """Pops the topmost frame from the stack"""
-        if len(self._frames) == 1:
-            raise IndexError
-        self._frames.pop()
-        self._current_frame = self._frames[-1]
+        self.__parent_frames.append(other)
+        other.__parent_frames.append(self)
+
+    def __setitem__(self, key: type[K], value: type[V]) -> None:
+        self.__current_frame[key] = value
 
     def __getitem__(self, key: type[K]) -> Optional[type[V]]:
-        for frame in reversed(self._frames):
-            if key in frame:
-                return frame[key]
+        return self.__getitem_helper(key, set())
+
+    def __getitem_helper(self, key: type[K], closed: set[str]) -> Optional[type[V]]:
+        if key in self.__current_frame:
+            return self.__current_frame[key]
+        for frame in self.__parent_frames:
+            if frame.__id in closed:
+                continue
+            closed.add(frame.__id)
+
+            found = frame.__getitem_helper(key, closed)
+            if found is not None:
+                return found
         return None
 
-    def __setitem__(self, key: type[K], value: type[V]) -> None:
-        self._current_frame[key] = value
+    def __iter__(self) -> Iterator[tuple[type[K], type[V]]]:
+        raise NotImplementedError
 
-    def __iter__(self):
-        return chain(*reversed(self._frames))
+    def __len__(self) -> int:
+        raise NotImplementedError
 
-    def __len__(self):
-        return len(self._frames)
 
+@dataclass(kw_only=True)
+class Node:
+    """
+    Represents actual data in an instance of client params. Unlike just the "param", a node has a
+    value and a parent.
+
+    To help with evaluating conds, it also keeps track of whether it is active/inactive and has a
+    context that can be used to resolve relative references to other nodes.
+    """
+
+    name: str = ""
+    "Param name or key name (depending on whether this is a map or message value)"
+    value: Any = None
+    parent: Optional["Node"] = None
+    children: list["Node"] = field(default_factory=list)
+
+    is_active: Optional[bool] = None
+    context: ContextFrame[str | int, "Node"] = field(default_factory=ContextFrame)
 
-class CondHelper:
-    @dataclass
-    class _Node:
-        """
-        Represents an actual field in an instance of client params. Unlike the "param" it is
-        an instance of, a node has a value, a parent, and can be active/inactive.
-        """
 
-        name: str
-        "Param name"
-        parent_name: Optional[str]
-        "Param name of parent"
-        value: Any
-        is_active: Optional[bool]
+class CondHelper:
 
     def __init__(self):
         self.params = ParamMap(SimulationParam)
-        self.node_by_name = ContextMap[str, self._Node]()
-        self.node_by_number = ContextMap[int, self._Node]()
+        self.nodes_by_surface: dict[str, list[Node]] = defaultdict(
+            list
+        )  # nodes with surface references
 
     def prune(self, tree: dict[str, Any]) -> dict[str, Any]:
         """
         Takes a dict of client parameters and returns a pruned tree, with inactive nodes and empty
         values removed.
         """
-        # Scan subtree to register all nodes.
-        self._scan_subtree(tree, None)
-        # Then, prune, using our registered nodes to evaluate conds.
-        return self._prune_subtree(tree)
-
-    def _scan_subtree(self, subtree: dict[str, Any], parent_name: Optional[str]) -> None:
+        # Scan entire tree to register all nodes.
+        root_node = Node(is_active=True)
+        self._scan_tree(tree, root_node)
+
+        physics_nodes_by_id = self._physics_nodes_by_id(root_node)
+        material_nodes_by_id = self._material_nodes_by_id(root_node)
+
+        # Entangle physics and material nodes
+        self._entangle_fluid_and_heat_physics(physics_nodes_by_id)
+        self._entangle_physics_and_material_nodes(tree, physics_nodes_by_id, material_nodes_by_id)
+        self._entangle_surface_nodes()
+
+        # Finally, prune, using our registered nodes to evaluate conds.
+        return self._prune_helper(root_node)
+
+    def _scan_tree(
+        self,
+        tree: dict[str, Any],
+        parent_node: Optional[Node],
+    ) -> None:
         """
-        From the root, traverses the tree to register all nodes, ignoring "repeated" and
-        "map" fields (we will only recurse into those fields after pushing a new frame onto the node
-        maps).
+        From the root, traverses the tree to register all nodes. Whenever we go down into a
+        "repeated" or "map" fields we push a new frame onto the context.
         """
-        for key, value in subtree.items():
+        context = parent_node.context
+        for key, value in tree.items():
             param = self.params.by_name.get(key, None)
             if param is None:
+                continue  # ignore non-params
+
+            node = Node(
+                name=key,
+                parent=parent_node,
+                value=value,
+                is_active=None,
+                context=context,
+            )
+            parent_node.children.append(node)
+            context[param.name] = node
+            context[param.number] = node
+
+            # Keep track of nodes with surface references
+            if is_surface_reference(param):
+                for surface in value:  # NOTE: just going to assume it's a list of surfaces
+                    self.nodes_by_surface[surface].append(parent_node)
+
+            if param.message_type is None or isinstance(value, str):
                 continue
-            node = self._Node(name=key, parent_name=parent_name, value=value, is_active=None)
 
-            if param.message_type is not None and param.label != FieldDescriptor.LABEL_REPEATED:
-                self._scan_subtree(value, param.name)
+            # Recurse into message types
+            if param.label == FieldDescriptor.LABEL_REPEATED:
+                # map types show up as REPEATED in descriptor
+                if isinstance(value, dict) and len(value) > 0:
+                    for k, v in value.items():
+                        new_node = Node(
+                            name=k,
+                            value=v,
+                            is_active=True,
+                            context=context.new_frame(),
+                        )  # new frame
+                        node.children.append(new_node)
+                        self._scan_tree(v, new_node)
+                elif isinstance(value, Iterable) and len(value) > 0:
+                    for v in value:
+                        new_node = Node(
+                            value=v,
+                            is_active=True,
+                            context=context.new_frame(),
+                        )  # new frame
+                        node.children.append(new_node)
+                        self._scan_tree(v, new_node)
+            else:
+                self._scan_tree(value, node)
+
+    def _entangle_fluid_and_heat_physics(self, physics_nodes_by_id: dict[str, Node]) -> None:
+        """
+        Entangles fluid and heat physics nodes.
+        """
+        fluid_node = None
+        heat_node = None
+        for node in physics_nodes_by_id.values():
+            if "fluid" in node.value:
+                fluid_node = node
+            elif "heat" in node.value:
+                heat_node = node
+        if fluid_node is not None and heat_node is not None:
+            fluid_node.context.entangle(heat_node.context)
+
+    def _entangle_physics_and_material_nodes(
+        self,
+        root: dict[str, Any],
+        physics_nodes_by_id: dict[str, Node],
+        material_nodes_by_id: dict[str, Node],
+    ) -> None:
+        """
+        Entangles physics and material nodes that have a relationship with the same volume.
+        """
+        for physics_node in physics_nodes_by_id.values():
+            physics_id = physics_node.value["physics_identifier"]["id"]
+            material_id = self._get_material_id_by_physics_id(root, physics_id)
+            if not material_id:
+                continue
+            material_node = material_nodes_by_id[material_id]
+            physics_node.context.entangle(material_node.context)
 
-            self.node_by_name[param.name] = node
-            self.node_by_number[param.number] = node
+    def _entangle_surface_nodes(self) -> None:
+        """
+        Entangles any nodes that reference the same surface.
+        """
+        for node_group in self.nodes_by_surface.values():
+            for a, b in pairwise(node_group):
+                a.context.entangle(b.context)
 
-    def _prune_subtree(self, subtree: dict[str, Any]) -> dict[str, Any]:
+    def _prune_helper(self, root_node: Node) -> dict[str, Any]:
         """
-        Prunes the tree, using the node maps to evaluate the conds.
+        Prunes the tree, using nodes to resolve conds.
         """
         pruned_tree = dict()
-        for key, value in subtree.items():
-            param = self.params.by_name.get(key, None)
-            if (
-                param is None or key == "surface_name"
-            ):  # "surface_name" is repeated twice in the tree...
-                pruned_tree[key] = value
+        for node in root_node.children:
+            name, value = node.name, node.value
+            param = self.params.by_name.get(name, None)
+
+            if param is None or name == "surface_name":  # "surface_name" appears twice in the tree
+                pruned_tree[name] = value
                 continue
-            node = self.node_by_name[key]
+
             if self._is_node_active(node):
                 if param.message_type is None or not str.startswith(
                     param.message_type.full_name, "luminary.proto.client"
                 ):
                     if not (isinstance(value, Iterable) and len(value) == 0):
-                        pruned_tree[key] = value
-                elif key == "physics":
-                    pruned_tree[key] = [self._prune_single_physics(subtree, v) for v in value]
+                        pruned_tree[name] = value
                 elif param.label == FieldDescriptor.LABEL_REPEATED:
                     # map types show up as REPEATED in descriptor, so we need to distinguish
                     if isinstance(value, dict) and len(value) > 0:
-                        pruned_tree[key] = {k: self._prune_repeated(v) for k, v in value.items()}
+                        pruned_tree[name] = {n.name: self._prune_helper(n) for n in node.children}
                     elif isinstance(value, Iterable) and len(value) > 0:
-                        pruned_tree[key] = [self._prune_repeated(v) for v in value]
+                        pruned_tree[name] = [self._prune_helper(n) for n in node.children]
                 else:
-                    pruned_value = self._prune_subtree(value)
-                    if len(pruned_value) > 0:
-                        pruned_tree[key] = pruned_value
+                    pruned_value = self._prune_helper(node)
+                    if len(pruned_value) > 0:  # prune empty dicts
+                        pruned_tree[name] = pruned_value
         return pruned_tree
 
-    def _prune_repeated(self, value: Any) -> dict[str, Any]:
-        """Prunes a single element of a repeated field."""
-        if not isinstance(value, dict):
-            return value
-        with self.node_by_name.scope():
-            with self.node_by_number.scope():
-                return self.prune(value)
-
-    def _prune_single_physics(
-        self, root: dict[str, Any], physics: dict[str, Any]
-    ) -> dict[str, Any]:
+    def _physics_nodes_by_id(self, root_node: Node) -> dict[str, Node]:
         """
-        Prunes a single physics object of the repeated physics field
-
-        Attempts to find the material associated with the physics, and if it finds one, scans the
-        material subtree to register all material properties in the node context before calling
-        prune() on the physics.
+        Gets a map of physics nodes by their id.
         """
-        try:
-            physics_id = physics["physics_identifier"]["id"]
-            material_id = self._get_material_id_by_physics_id(root, physics_id)
-            material = self._get_material_by_id(root, material_id)
-        except (KeyError, StopIteration, TypeError):
-            material = None
-        with self.node_by_name.scope():
-            with self.node_by_number.scope():
-                if material is not None:
-                    self._scan_subtree(material, None)
-                return self.prune(physics)
+        nodes_by_physics_id = dict()
+        for node in root_node.children:
+            if node.name == "physics":
+                for physics_node in node.children:
+                    physics_id = physics_node.value["physics_identifier"]["id"]
+                    nodes_by_physics_id[physics_id] = physics_node
+        return nodes_by_physics_id
+
+    def _material_nodes_by_id(self, root_node: Node) -> dict[str, Node]:
+        """
+        Gets a map of material nodes by their id.
+        """
+        nodes_by_material_id = dict()
+        for node in root_node.children:
+            if node.name == "material_entity":
+                for material_node in node.children:
+                    material_id = material_node.value["material_identifier"]["id"]
+                    nodes_by_material_id[material_id] = material_node
+        return nodes_by_material_id
 
     def _get_material_id_by_physics_id(
         self, root: dict[str, Any], physics_id: str
-    ) -> dict[str, Any]:
+    ) -> dict[str, Any] | None:
         volume_id = self._get_volume_id_by_physics_id(root, physics_id)
+        if volume_id is None:
+            return None
         return self._get_material_id_by_volume_id(root, volume_id)
 
-    def _get_volume_id_by_physics_id(self, root: dict[str, Any], physics_id: str) -> str:
-        entity_relationships = root["entity_relationships"]
-        return next(
-            link["volume_identifier"]["id"]
-            for link in entity_relationships["volume_physics_relationship"]
-            if link["physics_identifier"]["id"] == physics_id
-        )
-
-    def _get_material_id_by_volume_id(self, root: dict[str, Any], volume_id: str) -> str:
-        entity_relationships = root["entity_relationships"]
-        return next(
-            link["material_identifier"]["id"]
-            for link in entity_relationships["volume_material_relationship"]
-            if link["volume_identifier"]["id"] == volume_id
-        )
-
-    def _get_material_by_id(self, root: dict[str, Any], material_id: str) -> dict[str, Any]:
-        return next(
-            m for m in root["material_entity"] if m["material_identifier"]["id"] == material_id
-        )
-
-    def _is_node_active(self, node: "_Node") -> bool:
+    def _get_volume_id_by_physics_id(self, root: dict[str, Any], physics_id: str) -> str | None:
+        entity_relationships = root.get("entity_relationships")
+        if not entity_relationships:
+            return None
+        for link in entity_relationships.get("volume_physics_relationship", []):
+            if link.get("physics_identifier", {}).get("id") == physics_id:
+                return link.get("volume_identifier", {}).get("id")
+        return None
+
+    def _get_material_id_by_volume_id(self, root: dict[str, Any], volume_id: str) -> str | None:
+        entity_relationships = root.get("entity_relationships")
+        if not entity_relationships:
+            return None
+        for link in entity_relationships.get("volume_material_relationship", []):
+            if link.get("volume_identifier", {}).get("id") == volume_id:
+                return link.get("material_identifier", {}).get("id")
+        return None
+
+    def _is_node_active(self, node: Node) -> bool:
         """
         Checks if a node is active.
 
         Checks the parent first, and then the conds. Only does this once.
         """
         if node.is_active is None:
-            if node.parent_name is not None:
-                parent = self.node_by_name[node.parent_name]
-                if not self._is_node_active(parent):
+            if node.parent is not None:
+                if not self._is_node_active(node.parent):
                     node.is_active = False
                     return False
             cond = get_cond(self.params.by_name[node.name])
-            node.is_active = self._check_cond(cond)
+            node.is_active = self._check_cond(cond, node.context)
         return node.is_active
 
-    def _check_cond(self, cond: conditionpb.Condition) -> bool:
+    def _check_cond(
+        self, cond: conditionpb.Condition, context: ContextFrame[str | int, Node]
+    ) -> bool:
         type = cond.WhichOneof("typ")
         if type is None:
             return True
         if type == "choice":
-            return self._check_choice(cond.choice)
+            return self._check_choice(cond.choice, context)
         elif type == "boolean":
-            return self._check_boolean(cond.boolean)
+            return self._check_boolean(cond.boolean, context)
         elif type == "allof":
-            return all(self._check_cond(c) for c in cond.allof.cond)
+            return all(self._check_cond(c, context) for c in cond.allof.cond)
         elif type == "anyof":
-            return any(self._check_cond(c) for c in cond.anyof.cond)
+            return any(self._check_cond(c, context) for c in cond.anyof.cond)
         elif type == "not":
-            return not self._check_cond(getattr(cond, "not").cond)
+            return not self._check_cond(getattr(cond, "not").cond, context)
         elif type == "tag":
-            return self._check_tag(cond.tag)
+            return self._check_tag(cond.tag, context)
         elif type == "false":
             return (
                 True  # CondFalse is generally used to control visibility in UI and can be ignored
             )
         return False
 
-    def _check_choice(self, cond_choice: conditionpb.Choice) -> bool:
+    def _check_choice(
+        self, cond_choice: conditionpb.Choice, context: ContextFrame[str | int, Node]
+    ) -> bool:
         name = cond_choice.param_name
-        node = self.node_by_name[name]
+        node = context[name]
         if node is None:
-            param = self.params.by_name[name]
-            return param == cond_choice.tag
+            return 0 == cond_choice.tag
         return self._is_node_active(node) and node.value == cond_choice.name
 
-    def _check_boolean(self, cond_boolean: conditionpb.TrueFalse) -> bool:
+    def _check_boolean(
+        self, cond_boolean: conditionpb.TrueFalse, context: ContextFrame[str | int, Node]
+    ) -> bool:
         tag = cond_boolean.param_name_tag
-        node = self.node_by_number[tag]
+        node = context[tag]
         if node is None:
-            param = self.params.by_number[tag]
-            return bool(param)
+            return False
         return self._is_node_active(node) and bool(node.value)
 
-    def _check_tag(self, cond_tag: conditionpb.Tag) -> bool:
+    def _check_tag(self, cond_tag: conditionpb.Tag, context: ContextFrame[str | int, Node]) -> bool:
         name = cond_tag.tag_name
-        node = self.node_by_name[name]
-        return node is not None and self._is_node_active(node)
+        node = context[name]
+        return (node is not None) and self._is_node_active(node)
 
 
 def params_to_dict(sim_params: SimulationParam) -> dict[str, Any]:
@@ -340,3 +433,7 @@ def get_default(param: FieldDescriptor) -> Any:
     elif type == "vector3":
         return dfl.vector3
     return None
+
+
+def is_surface_reference(param: FieldDescriptor) -> bool:
+    return param.GetOptions().Extensions[optionspb.is_geometry_surfaces]

luminarycloud/_helpers/defaults.py

@@ -1,13 +1,13 @@
 # Copyright 2024 Luminary Cloud, Inc. All Rights Reserved.
-from google.protobuf.any_pb2 import Any as AnyPB
 from google.protobuf.descriptor import FieldDescriptor
+from google.protobuf.message import Message
 
 import luminarycloud._proto.base.base_pb2 as basepb
 import luminarycloud._proto.options.options_pb2 as optionspb
 from luminarycloud.types.vector3 import _to_vector3_proto
 
 
-def _reset_defaults(params: AnyPB) -> None:
+def _reset_defaults(params: Message) -> None:
     for field in params.DESCRIPTOR.fields:
         if field.label == FieldDescriptor.LABEL_REPEATED:
             for nested in getattr(params, field.name):
@@ -18,7 +18,7 @@ def _reset_defaults(params: AnyPB) -> None:
                 _reset_defaults(getattr(params, field.name))
 
 
-def _set_default(params: AnyPB, name: str) -> None:
+def _set_default(params: Message, name: str) -> None:
     field = params.DESCRIPTOR.fields_by_name[name]
     dfl: optionspb.Value = field.GetOptions().Extensions[optionspb.default_value]
     type = dfl.WhichOneof("typ")

luminarycloud/_helpers/download.py

@@ -96,12 +96,23 @@ def download_surface_solution(
 def download_volume_solution(
     client: Client,
     solution_id: str,
+    single_precision: bool = False,
 ) -> FileChunkStream:
     """
     Returns the download as a file-like object.
 
     The filename can be retrieved from the `filename` attribute of the returned object.
 
+    Parameters
+    ----------
+    client: Client
+        The client to use for the download.
+    solution_id: str
+        The ID of the solution to download.
+    single_precision: bool
+        If True, the solution will be downloaded in single precision.
+        If False, the solution will be downloaded in double precision.
+
     Examples
     --------
     >>> with download_volume_solution(client, "your-solution-id") as dl:
@@ -109,7 +120,7 @@ def download_volume_solution(
     ...         fp.write(dl.read())
     """
 
-    request = GetSolutionVolumeDataRequest(id=solution_id)
+    request = GetSolutionVolumeDataRequest(id=solution_id, single_precision=single_precision)
     response = client.GetSolutionVolumeData(request)
     return _create_file_chunk_stream(client, solution_id, "volume solution", response.file)