opentrons 8.3.1a1__py2.py3-none-any.whl → 8.4.0a1__py2.py3-none-any.whl

opentrons/protocol_engine/protocol_engine.py

@@ -1,29 +1,32 @@
 """ProtocolEngine class definition."""
+
 from contextlib import AsyncExitStack
 from logging import getLogger
 from typing import Dict, Optional, Union, AsyncGenerator, Callable
-from opentrons.protocol_engine.actions.actions import (
-    ResumeFromRecoveryAction,
-    SetErrorRecoveryPolicyAction,
-)
 
-from opentrons.protocols.models import LabwareDefinition
-from opentrons.hardware_control import HardwareControlAPI
-from opentrons.hardware_control.modules import AbstractModule as HardwareModuleAPI
-from opentrons.hardware_control.types import PauseType as HardwarePauseType
 from opentrons_shared_data.errors import (
     ErrorCodes,
     EnumeratedError,
 )
+from opentrons_shared_data.labware.labware_definition import LabwareDefinition
 
+from opentrons.hardware_control import HardwareControlAPI
+from opentrons.hardware_control.modules import AbstractModule as HardwareModuleAPI
+from opentrons.hardware_control.types import PauseType as HardwarePauseType
+
+from .actions.actions import (
+    ResumeFromRecoveryAction,
+    SetErrorRecoveryPolicyAction,
+)
 from .errors import ProtocolCommandFailedError, ErrorOccurrence, CommandNotAllowedError
 from .errors.exceptions import EStopActivatedError
 from .error_recovery_policy import ErrorRecoveryPolicy
-from . import commands, slot_standardization
+from . import commands, slot_standardization, labware_offset_standardization
 from .resources import ModelUtils, ModuleDataProvider, FileProvider
 from .types import (
     LabwareOffset,
     LabwareOffsetCreate,
+    LegacyLabwareOffsetCreate,
     LabwareUri,
     ModuleModel,
     Liquid,
@@ -516,15 +519,21 @@ class ProtocolEngine:
             )
         )
 
-    def add_labware_offset(self, request: LabwareOffsetCreate) -> LabwareOffset:
+    def add_labware_offset(
+        self, request: LabwareOffsetCreate | LegacyLabwareOffsetCreate
+    ) -> LabwareOffset:
         """Add a new labware offset and return it.
 
         The added offset will apply to subsequent `LoadLabwareCommand`s.
 
         To retrieve offsets later, see `.state_view.labware`.
         """
-        request = slot_standardization.standardize_labware_offset(
-            request, self.state_view.config.robot_type
+        internal_request = (
+            labware_offset_standardization.standardize_labware_offset_create(
+                request,
+                self.state_view.config.robot_type,
+                self.state_view.addressable_areas.deck_definition,
+            )
         )
 
         labware_offset_id = self._model_utils.generate_id()
@@ -533,7 +542,7 @@ class ProtocolEngine:
             AddLabwareOffsetAction(
                 labware_offset_id=labware_offset_id,
                 created_at=created_at,
-                request=request,
+                request=internal_request,
             )
         )
         return self.state_view.labware.get_labware_offset(

opentrons/protocol_engine/resources/deck_configuration_provider.py

@@ -1,7 +1,11 @@
 """Deck configuration resource provider."""
+
 from typing import List, Set, Tuple
 
-from opentrons_shared_data.deck.types import DeckDefinitionV5, CutoutFixture
+from opentrons_shared_data.deck.types import (
+    DeckDefinitionV5,
+    CutoutFixture,
+)
 
 from opentrons.types import DeckSlotName
 
@@ -17,6 +21,7 @@ from ..errors import (
     CutoutDoesNotExistError,
     FixtureDoesNotExistError,
     AddressableAreaDoesNotExistError,
+    SlotDoesNotExistError,
 )
 
 
@@ -98,12 +103,15 @@ def get_potential_cutout_fixtures(
 def get_addressable_area_from_name(
     addressable_area_name: str,
     cutout_position: DeckPoint,
-    base_slot: DeckSlotName,
     deck_definition: DeckDefinitionV5,
 ) -> AddressableArea:
     """Given a name and a cutout position, get an addressable area on the deck."""
     for addressable_area in deck_definition["locations"]["addressableAreas"]:
         if addressable_area["id"] == addressable_area_name:
+            cutout_id, _ = get_potential_cutout_fixtures(
+                addressable_area_name, deck_definition
+            )
+            base_slot = get_deck_slot_for_cutout_id(cutout_id)
             area_offset = addressable_area["offsetFromCutoutFixture"]
             position = AddressableOffsetVector(
                 x=area_offset[0] + cutout_position.x,
@@ -130,3 +138,87 @@ def get_addressable_area_from_name(
     raise AddressableAreaDoesNotExistError(
         f"Could not find addressable area with name {addressable_area_name}"
     )
+
+
+def get_deck_slot_for_cutout_id(cutout_id: str) -> DeckSlotName:
+    """Get the corresponding deck slot for an addressable area."""
+    try:
+        return CUTOUT_TO_DECK_SLOT_MAP[cutout_id]
+    except KeyError:
+        raise CutoutDoesNotExistError(f"Could not find data for cutout {cutout_id}")
+
+
+def get_cutout_id_by_deck_slot_name(slot_name: DeckSlotName) -> str:
+    """Get the Cutout ID of a given Deck Slot by Deck Slot Name."""
+    try:
+        return DECK_SLOT_TO_CUTOUT_MAP[slot_name]
+    except KeyError:
+        raise SlotDoesNotExistError(f"Could not find data for slot {slot_name.value}")
+
+
+def get_labware_hosting_addressable_area_name_for_cutout_and_cutout_fixture(
+    cutout_id: str, cutout_fixture_id: str, deck_definition: DeckDefinitionV5
+) -> str:
+    """Get the first addressable area that can contain labware for a cutout and fixture.
+
+    This probably isn't relevant outside of labware offset locations, where (for now) nothing
+    provides more than one labware-containing addressable area.
+    """
+    for cutoutFixture in deck_definition["cutoutFixtures"]:
+        if cutoutFixture["id"] != cutout_fixture_id:
+            continue
+        provided_aas = cutoutFixture["providesAddressableAreas"].get(cutout_id, None)
+        if provided_aas is None:
+            raise CutoutDoesNotExistError(
+                f"{cutout_fixture_id} does not go in {cutout_id}"
+            )
+        for aa_id in provided_aas:
+            for addressable_area in deck_definition["locations"]["addressableAreas"]:
+                if addressable_area["id"] != aa_id:
+                    continue
+                # TODO: In deck def v6 this will be easier, but as of right now there isn't really
+                # a way to tell from an addressable area whether it takes labware so let's take the
+                # first one
+                return aa_id
+            raise AddressableAreaDoesNotExistError(
+                f"Could not find an addressable area that allows labware from cutout fixture {cutout_fixture_id} in cutout {cutout_id}"
+            )
+
+    raise FixtureDoesNotExistError(f"Could not find entry for {cutout_fixture_id}")
+
+
+# This is a temporary shim while Protocol Engine's conflict-checking code
+# can only take deck slots as input.
+# Long-term solution: Check for conflicts based on bounding boxes, not slot adjacencies.
+# Shorter-term: Change the conflict-checking code to take cutouts instead of deck slots.
+CUTOUT_TO_DECK_SLOT_MAP: dict[str, DeckSlotName] = {
+    # OT-2
+    "cutout1": DeckSlotName.SLOT_1,
+    "cutout2": DeckSlotName.SLOT_2,
+    "cutout3": DeckSlotName.SLOT_3,
+    "cutout4": DeckSlotName.SLOT_4,
+    "cutout5": DeckSlotName.SLOT_5,
+    "cutout6": DeckSlotName.SLOT_6,
+    "cutout7": DeckSlotName.SLOT_7,
+    "cutout8": DeckSlotName.SLOT_8,
+    "cutout9": DeckSlotName.SLOT_9,
+    "cutout10": DeckSlotName.SLOT_10,
+    "cutout11": DeckSlotName.SLOT_11,
+    "cutout12": DeckSlotName.FIXED_TRASH,
+    # Flex
+    "cutoutA1": DeckSlotName.SLOT_A1,
+    "cutoutA2": DeckSlotName.SLOT_A2,
+    "cutoutA3": DeckSlotName.SLOT_A3,
+    "cutoutB1": DeckSlotName.SLOT_B1,
+    "cutoutB2": DeckSlotName.SLOT_B2,
+    "cutoutB3": DeckSlotName.SLOT_B3,
+    "cutoutC1": DeckSlotName.SLOT_C1,
+    "cutoutC2": DeckSlotName.SLOT_C2,
+    "cutoutC3": DeckSlotName.SLOT_C3,
+    "cutoutD1": DeckSlotName.SLOT_D1,
+    "cutoutD2": DeckSlotName.SLOT_D2,
+    "cutoutD3": DeckSlotName.SLOT_D3,
+}
+DECK_SLOT_TO_CUTOUT_MAP = {
+    deck_slot: cutout for cutout, deck_slot in CUTOUT_TO_DECK_SLOT_MAP.items()
+}

opentrons/protocol_engine/resources/deck_data_provider.py

@@ -10,7 +10,7 @@ from opentrons_shared_data.deck import (
     DEFAULT_DECK_DEFINITION_VERSION,
 )
 from opentrons_shared_data.deck.types import DeckDefinitionV5
-from opentrons.protocols.models import LabwareDefinition
+from opentrons_shared_data.labware.labware_definition import LabwareDefinition
 from opentrons.types import DeckSlotName
 
 from ..types import (

opentrons/protocol_engine/resources/labware_data_provider.py

@@ -6,7 +6,12 @@ abstract away rough edges until we can improve those underlying interfaces.
 import logging
 from anyio import to_thread
 
-from opentrons.protocols.models import LabwareDefinition
+from opentrons_shared_data.labware.labware_definition import (
+    LabwareDefinition,
+    LabwareDefinition3,
+    labware_definition_type_adapter,
+)
+
 from opentrons.protocols.labware import get_labware_definition
 
 # TODO (lc 09-26-2022) We should conditionally import ot2 or ot3 calibration
@@ -44,7 +49,7 @@ class LabwareDataProvider:
     def _get_labware_definition_sync(
         load_name: str, namespace: str, version: int
     ) -> LabwareDefinition:
-        return LabwareDefinition.model_validate(
+        return labware_definition_type_adapter.validate_python(
             get_labware_definition(load_name, namespace, version)
         )
 
@@ -72,15 +77,30 @@ class LabwareDataProvider:
         labware_definition: LabwareDefinition,
         nominal_fallback: float,
     ) -> float:
-        try:
-            return instr_cal.load_tip_length_for_pipette(
-                pipette_serial, labware_definition
-            ).tip_length
-
-        except TipLengthCalNotFound as e:
-            message = (
-                f"No calibrated tip length found for {pipette_serial},"
-                f" using nominal fallback value of {nominal_fallback}"
+        if isinstance(labware_definition, LabwareDefinition3):
+            # FIXME(mm, 2025-02-19): This needs to be resolved for v8.4.0.
+            # Tip length calibration internals don't yet support schema 3 because
+            # it's probably an incompatible change at the filesystem level
+            # (not downgrade-safe), and because robot-server's calibration sessions
+            # are built atop opentrons.protocol_api.core.legacy, which does not (yet?)
+            # support labware schema 3.
+            # https://opentrons.atlassian.net/browse/EXEC-1230
+            log.warning(
+                f"Tip rack"
+                f" {labware_definition.namespace}/{labware_definition.parameters.loadName}/{labware_definition.version}"
+                f" has schema 3, so tip length calibration is currently unsupported."
+                f" Using nominal fallback of {nominal_fallback}."
             )
-            log.debug(message, exc_info=e)
             return nominal_fallback
+        else:
+            try:
+                return instr_cal.load_tip_length_for_pipette(
+                    pipette_serial, labware_definition
+                ).tip_length
+            except TipLengthCalNotFound as e:
+                message = (
+                    f"No calibrated tip length found for {pipette_serial},"
+                    f" using nominal fallback value of {nominal_fallback}"
+                )
+                log.debug(message, exc_info=e)
+                return nominal_fallback

opentrons/protocol_engine/resources/labware_validation.py

@@ -1,7 +1,9 @@
 """Validation file for labware role and location checking functions."""
 
-from opentrons_shared_data.labware.labware_definition import LabwareRole
-from opentrons.protocols.models import LabwareDefinition
+from opentrons_shared_data.labware.labware_definition import (
+    LabwareDefinition,
+    LabwareRole,
+)
 
 
 def is_flex_trash(load_name: str) -> bool:
@@ -14,9 +16,9 @@ def is_absorbance_reader_lid(load_name: str) -> bool:
     return load_name == "opentrons_flex_lid_absorbance_plate_reader_module"
 
 
-def is_evotips(load_name: str) -> bool:
-    """Check if a labware is an evotips tiprack."""
-    return load_name == "evotips_opentrons_96_labware"
+def is_lid_stack(load_name: str) -> bool:
+    """Check if a labware object is a system lid stack object."""
+    return load_name == "protocol_engine_lid_stack_object"
 
 
 def validate_definition_is_labware(definition: LabwareDefinition) -> bool:

opentrons/protocol_engine/slot_standardization.py

@@ -14,7 +14,6 @@ This module does that conversion, for any Protocol Engine input that contains a
 deck slot.
 """
 
-
 from typing import Any, Callable, Dict, Type
 
 from opentrons_shared_data.robot.types import RobotType
@@ -22,32 +21,15 @@ from opentrons_shared_data.robot.types import RobotType
 from . import commands
 from .types import (
     OFF_DECK_LOCATION,
+    SYSTEM_LOCATION,
     DeckSlotLocation,
-    LabwareLocation,
+    LoadableLabwareLocation,
     AddressableAreaLocation,
-    LabwareOffsetCreate,
     ModuleLocation,
     OnLabwareLocation,
 )
 
 
-def standardize_labware_offset(
-    original: LabwareOffsetCreate, robot_type: RobotType
-) -> LabwareOffsetCreate:
-    """Convert the deck slot in the given `LabwareOffsetCreate` to match the given robot type."""
-    return original.model_copy(
-        update={
-            "location": original.location.model_copy(
-                update={
-                    "slotName": original.location.slotName.to_equivalent_for_robot_type(
-                        robot_type
-                    )
-                }
-            )
-        }
-    )
-
-
 def standardize_command(
     original: commands.CommandCreate, robot_type: RobotType
 ) -> commands.CommandCreate:
@@ -119,15 +101,21 @@ _standardize_command_functions: Dict[
 
 
 def _standardize_labware_location(
-    original: LabwareLocation, robot_type: RobotType
-) -> LabwareLocation:
+    original: LoadableLabwareLocation, robot_type: RobotType
+) -> LoadableLabwareLocation:
     if isinstance(original, DeckSlotLocation):
         return _standardize_deck_slot_location(original, robot_type)
     elif (
         isinstance(
-            original, (ModuleLocation, OnLabwareLocation, AddressableAreaLocation)
+            original,
+            (
+                ModuleLocation,
+                OnLabwareLocation,
+                AddressableAreaLocation,
+            ),
         )
         or original == OFF_DECK_LOCATION
+        or original == SYSTEM_LOCATION
     ):
         return original
 

opentrons/protocol_engine/state/addressable_areas.py

@@ -1,4 +1,5 @@
 """Basic addressable area data state and store."""
+
 from dataclasses import dataclass
 from functools import cached_property
 from typing import Dict, List, Optional, Set
@@ -112,43 +113,6 @@ def _get_conflicting_addressable_areas_error_string(
     return ", ".join(display_names)
 
 
-# This is a temporary shim while Protocol Engine's conflict-checking code
-# can only take deck slots as input.
-# Long-term solution: Check for conflicts based on bounding boxes, not slot adjacencies.
-# Shorter-term: Change the conflict-checking code to take cutouts instead of deck slots.
-CUTOUT_TO_DECK_SLOT_MAP: Dict[str, DeckSlotName] = {
-    # OT-2
-    "cutout1": DeckSlotName.SLOT_1,
-    "cutout2": DeckSlotName.SLOT_2,
-    "cutout3": DeckSlotName.SLOT_3,
-    "cutout4": DeckSlotName.SLOT_4,
-    "cutout5": DeckSlotName.SLOT_5,
-    "cutout6": DeckSlotName.SLOT_6,
-    "cutout7": DeckSlotName.SLOT_7,
-    "cutout8": DeckSlotName.SLOT_8,
-    "cutout9": DeckSlotName.SLOT_9,
-    "cutout10": DeckSlotName.SLOT_10,
-    "cutout11": DeckSlotName.SLOT_11,
-    "cutout12": DeckSlotName.FIXED_TRASH,
-    # Flex
-    "cutoutA1": DeckSlotName.SLOT_A1,
-    "cutoutA2": DeckSlotName.SLOT_A2,
-    "cutoutA3": DeckSlotName.SLOT_A3,
-    "cutoutB1": DeckSlotName.SLOT_B1,
-    "cutoutB2": DeckSlotName.SLOT_B2,
-    "cutoutB3": DeckSlotName.SLOT_B3,
-    "cutoutC1": DeckSlotName.SLOT_C1,
-    "cutoutC2": DeckSlotName.SLOT_C2,
-    "cutoutC3": DeckSlotName.SLOT_C3,
-    "cutoutD1": DeckSlotName.SLOT_D1,
-    "cutoutD2": DeckSlotName.SLOT_D2,
-    "cutoutD3": DeckSlotName.SLOT_D3,
-}
-DECK_SLOT_TO_CUTOUT_MAP = {
-    deck_slot: cutout for cutout, deck_slot in CUTOUT_TO_DECK_SLOT_MAP.items()
-}
-
-
 class AddressableAreaStore(HasState[AddressableAreaState], HandlesActions):
     """Addressable area state container."""
 
@@ -221,13 +185,11 @@ class AddressableAreaStore(HasState[AddressableAreaState], HandlesActions):
             cutout_position = deck_configuration_provider.get_cutout_position(
                 cutout_id, deck_definition
             )
-            base_slot = CUTOUT_TO_DECK_SLOT_MAP[cutout_id]
             for addressable_area_name in provided_addressable_areas:
                 addressable_areas.append(
                     deck_configuration_provider.get_addressable_area_from_name(
                         addressable_area_name=addressable_area_name,
                         cutout_position=cutout_position,
-                        base_slot=base_slot,
                         deck_definition=deck_definition,
                     )
                 )
@@ -242,12 +204,10 @@ class AddressableAreaStore(HasState[AddressableAreaState], HandlesActions):
             cutout_position = deck_configuration_provider.get_cutout_position(
                 cutout_id, self._state.deck_definition
             )
-            base_slot = CUTOUT_TO_DECK_SLOT_MAP[cutout_id]
             addressable_area = (
                 deck_configuration_provider.get_addressable_area_from_name(
                     addressable_area_name=addressable_area_name,
                     cutout_position=cutout_position,
-                    base_slot=base_slot,
                     deck_definition=self._state.deck_definition,
                 )
             )
@@ -300,6 +260,11 @@ class AddressableAreaView:
         """
         self._state = state
 
+    @cached_property
+    def deck_definition(self) -> DeckDefinitionV5:
+        """The full deck definition."""
+        return self._state.deck_definition
+
     @cached_property
     def deck_extents(self) -> Point:
         """The maximum space on the deck."""
@@ -426,11 +391,9 @@ class AddressableAreaView:
         cutout_position = deck_configuration_provider.get_cutout_position(
             cutout_id, self._state.deck_definition
         )
-        base_slot = CUTOUT_TO_DECK_SLOT_MAP[cutout_id]
         return deck_configuration_provider.get_addressable_area_from_name(
             addressable_area_name=addressable_area_name,
             cutout_position=cutout_position,
-            base_slot=base_slot,
             deck_definition=self._state.deck_definition,
         )
 
@@ -526,7 +489,7 @@ class AddressableAreaView:
 
     def get_cutout_id_by_deck_slot_name(self, slot_name: DeckSlotName) -> str:
         """Get the Cutout ID of a given Deck Slot by Deck Slot Name."""
-        return DECK_SLOT_TO_CUTOUT_MAP[slot_name]
+        return deck_configuration_provider.get_cutout_id_by_deck_slot_name(slot_name)
 
     def get_fixture_by_deck_slot_name(
         self, slot_name: DeckSlotName
@@ -534,7 +497,9 @@ class AddressableAreaView:
         """Get the Cutout Fixture currently loaded where a specific Deck Slot would be."""
         deck_config = self._state.deck_configuration
         if deck_config:
-            slot_cutout_id = DECK_SLOT_TO_CUTOUT_MAP[slot_name]
+            slot_cutout_id = (
+                deck_configuration_provider.get_cutout_id_by_deck_slot_name(slot_name)
+            )
             slot_cutout_fixture = None
             # This will only ever be one under current assumptions
             for (
@@ -571,7 +536,9 @@ class AddressableAreaView:
         """Get the serial number provided by the deck configuration for a Fixture at a given location."""
         deck_config = self._state.deck_configuration
         if deck_config:
-            slot_cutout_id = DECK_SLOT_TO_CUTOUT_MAP[slot_name]
+            slot_cutout_id = (
+                deck_configuration_provider.get_cutout_id_by_deck_slot_name(slot_name)
+            )
             # This will only ever be one under current assumptions
             for (
                 cutout_id,
@@ -582,6 +549,44 @@ class AddressableAreaView:
                     return opentrons_module_serial_number
         return None
 
+    def get_serial_number_by_cutout_id(self, slot_cutout_id: str) -> str | None:
+        """Gets serial number from deck at a given cutout ID if one exists."""
+        deck_config = self._state.deck_configuration
+        if deck_config:
+            for (
+                cutout_id,
+                cutout_fixture_id,
+                opentrons_module_serial_number,
+            ) in deck_config:
+                if cutout_id == slot_cutout_id:
+                    return opentrons_module_serial_number
+        return None
+
+    def get_fixture_serial_from_deck_configuration_by_addressable_area(
+        self, addressable_area_name: str
+    ) -> Optional[str]:
+        """Get the serial number provided by the deck configuration for a Fixture that provides a given addressable area."""
+        deck_config = self._state.deck_configuration
+        if deck_config:
+            potential_fixtures = (
+                deck_configuration_provider.get_potential_cutout_fixtures(
+                    addressable_area_name, self._state.deck_definition
+                )
+            )
+            slot_cutout_id = potential_fixtures[0]
+            fixture_ids = [
+                fixture.cutout_fixture_id for fixture in potential_fixtures[1]
+            ]
+            # This will only ever be one under current assumptions
+            for (
+                cutout_id,
+                cutout_fixture_id,
+                opentrons_module_serial_number,
+            ) in deck_config:
+                if cutout_id == slot_cutout_id and cutout_fixture_id in fixture_ids:
+                    return opentrons_module_serial_number
+        return None
+
     def get_slot_definition(self, slot_id: str) -> SlotDefV3:
         """Get the definition of a slot in the deck.
 
@@ -658,3 +663,36 @@ class AddressableAreaView:
                 raise AreaNotInDeckConfigurationError(
                     f"{addressable_area_name} not provided by deck configuration."
                 )
+
+    def get_current_potential_cutout_fixtures_for_addressable_area(
+        self, addressable_area_name: str
+    ) -> tuple[str, Set[PotentialCutoutFixture]]:
+        """Get the set of cutout fixtures that might provide a given addressable area.
+
+        This takes into account the constraints already established by load commands or by a loaded deck
+        configuration, and may therefore return different results for the same addressable area at
+        different points in the protocol after deck configuration constraints have changed.
+
+        This returns the common cutout id and the potential fixtures.
+        """
+        (
+            cutout_id,
+            base_potential_fixtures,
+        ) = deck_configuration_provider.get_potential_cutout_fixtures(
+            addressable_area_name, self._state.deck_definition
+        )
+        try:
+            loaded_potential_fixtures = (
+                self._state.potential_cutout_fixtures_by_cutout_id[cutout_id]
+            )
+            return cutout_id, loaded_potential_fixtures.intersection(
+                base_potential_fixtures
+            )
+        except KeyError:
+            # If there was a key error here, it's because this function was (eventually) called
+            # from the body of a command implementation whose state update will load the
+            # addressable area it's querying... but that state update has not been submitted
+            # and processed, so nothing has created the entry for this cutout id yet. Do what
+            # we'll do when we actually get to that state update, which is apply the base
+            # potential fixtures from the deck def.
+            return cutout_id, base_potential_fixtures

opentrons/protocol_engine/state/frustum_helpers.py

@@ -5,6 +5,10 @@ from math import isclose
 
 from ..errors.exceptions import InvalidLiquidHeightFound
 
+from opentrons.protocol_engine.types.liquid_level_detection import (
+    LiquidTrackingType,
+    SimulatedProbeResult,
+)
 from opentrons_shared_data.labware.labware_definition import (
     InnerWellGeometry,
     WellSegment,
@@ -241,9 +245,8 @@ def _get_segment_capacity(segment: WellSegment) -> float:
 def get_well_volumetric_capacity(
     well_geometry: InnerWellGeometry,
 ) -> List[Tuple[float, float]]:
-    """Return the total volumetric capacity of a well as a map of height borders to volume."""
-    # dictionary map of heights to volumetric capacities within their respective segment
-    # {top_height_0: volume_0, top_height_1: volume_1, top_height_2: volume_2}
+    """Return the volumetric capacity of a well as a list of pairs relating segment heights to volumes."""
+    #  [(top_height_0, section_0_volume), (top_height_1, section_1_volume), ...]
     well_volume = []
 
     # get the well segments sorted in ascending order
@@ -341,7 +344,7 @@ def _find_volume_in_partial_frustum(
 ) -> float:
     """Look through a sorted list of frusta for a target height, and find the volume at that height."""
     for segment in sorted_well:
-        if segment.bottomHeight < target_height < segment.topHeight:
+        if segment.bottomHeight <= target_height <= segment.topHeight:
             relative_target_height = target_height - segment.bottomHeight
             section_height = segment.topHeight - segment.bottomHeight
             return volume_at_height_within_section(
@@ -356,9 +359,13 @@ def _find_volume_in_partial_frustum(
 
 
 def find_volume_at_well_height(
-    target_height: float, well_geometry: InnerWellGeometry
-) -> float:
+    target_height: LiquidTrackingType,
+    well_geometry: InnerWellGeometry,
+) -> LiquidTrackingType:
     """Find the volume within a well, at a known height."""
+    # comparisons with SimulatedProbeResult objects aren't meaningful, just return
+    if isinstance(target_height, SimulatedProbeResult):
+        return target_height
     volumetric_capacity = get_well_volumetric_capacity(well_geometry)
     max_height = volumetric_capacity[-1][0]
     if target_height < 0 or target_height > max_height:
@@ -417,13 +424,22 @@ def _find_height_in_partial_frustum(
 
 
 def find_height_at_well_volume(
-    target_volume: float, well_geometry: InnerWellGeometry
-) -> float:
+    target_volume: LiquidTrackingType,
+    well_geometry: InnerWellGeometry,
+    raise_error_if_result_invalid: bool = True,
+) -> LiquidTrackingType:
     """Find the height within a well, at a known volume."""
+    # comparisons with SimulatedProbeResult objects aren't meaningful, just
+    # return if we have one of those
+    if isinstance(target_volume, SimulatedProbeResult):
+        return target_volume
+
     volumetric_capacity = get_well_volumetric_capacity(well_geometry)
     max_volume = sum(row[1] for row in volumetric_capacity)
-    if target_volume < 0 or target_volume > max_volume:
-        raise InvalidLiquidHeightFound("Invalid target volume.")
+
+    if raise_error_if_result_invalid:
+        if target_volume < 0 or target_volume > max_volume:
+            raise InvalidLiquidHeightFound("Invalid target volume.")
 
     sorted_well = sorted(well_geometry.sections, key=lambda section: section.topHeight)
     # find the section the target volume is in and compute the height