feat(shared-data): add transitional shape calculations (#16554)

<!--
Thanks for taking the time to open a Pull Request (PR)! Please make sure
you've read the "Opening Pull Requests" section of our Contributing
Guide:


https://github.com/Opentrons/opentrons/blob/edge/CONTRIBUTING.md#opening-pull-requests

GitHub provides robust markdown to format your PR. Links, diagrams,
pictures, and videos along with text formatting make it possible to
create a rich and informative PR. For more information on GitHub
markdown, see:


https://docs.github.com/en/get-started/writing-on-github/getting-started-with-writing-and-formatting-on-github/basic-writing-and-formatting-syntax

To ensure your code is reviewed quickly and thoroughly, please fill out
the sections below to the best of your ability!
-->

# Overview

Adds a cached property to the "Squared cone" labware geometry segment
that computes the height/volume tables if needed.
Also update frustum helpers to use these and don't raise a Not
implemented error on this shape.


<!--
Describe your PR at a high level. State acceptance criteria and how this
PR fits into other work. Link issues, PRs, and other relevant resources.
-->

## Test Plan and Hands on Testing

<!--
Describe your testing of the PR. Emphasize testing not reflected in the
code. Attach protocols, logs, screenshots and any other assets that
support your testing.
-->

## Changelog

<!--
List changes introduced by this PR considering future developers and the
end user. Give careful thought and clear documentation to breaking
changes.
-->

## Review requests

<!--
- What do you need from reviewers to feel confident this PR is ready to
merge?
- Ask questions.
-->

## Risk assessment

<!--
- Indicate the level of attention this PR needs.
- Provide context to guide reviewers.
- Discuss trade-offs, coupling, and side effects.
- Look for the possibility, even if you think it's small, that your
change may affect some other part of the system.
- For instance, changing return tip behavior may also change the
behavior of labware calibration.
- How do your unit tests and on hands on testing mitigate this PR's
risks and the risk of future regressions?
- Especially in high risk PRs, explain how you know your testing is
enough.
-->

api/src/opentrons/protocol_engine/state/frustum_helpers.py

@@ -11,6 +11,7 @@
     SphericalSegment,
     ConicalFrustum,
     CuboidalFrustum,
+    SquaredConeSegment,
 )
 
 
@@ -127,6 +128,15 @@ def _volume_from_height_spherical(
     return volume
 
 
+def _volume_from_height_squared_cone(
+    target_height: float, segment: SquaredConeSegment
+) -> float:
+    """Find the volume given a height within a squared cone segment."""
+    heights = segment.height_to_volume_table.keys()
+    best_fit_height = min(heights, key=lambda x: abs(x - target_height))
+    return segment.height_to_volume_table[best_fit_height]
+
+
 def _height_from_volume_circular(
     volume: float,
     total_frustum_height: float,
@@ -197,15 +207,24 @@ def _height_from_volume_spherical(
     return height
 
 
+def _height_from_volume_squared_cone(
+    target_volume: float, segment: SquaredConeSegment
+) -> float:
+    """Find the height given a volume within a squared cone segment."""
+    volumes = segment.volume_to_height_table.keys()
+    best_fit_volume = min(volumes, key=lambda x: abs(x - target_volume))
+    return segment.volume_to_height_table[best_fit_volume]
+
+
 def _get_segment_capacity(segment: WellSegment) -> float:
+    section_height = segment.topHeight - segment.bottomHeight
     match segment:
         case SphericalSegment():
             return _volume_from_height_spherical(
                 target_height=segment.topHeight,
                 radius_of_curvature=segment.radiusOfCurvature,
             )
         case CuboidalFrustum():
-            section_height = segment.topHeight - segment.bottomHeight
             return _volume_from_height_rectangular(
                 target_height=section_height,
                 bottom_length=segment.bottomYDimension,
@@ -215,13 +234,14 @@ def _get_segment_capacity(segment: WellSegment) -> float:
                 total_frustum_height=section_height,
             )
         case ConicalFrustum():
-            section_height = segment.topHeight - segment.bottomHeight
             return _volume_from_height_circular(
                 target_height=section_height,
                 total_frustum_height=section_height,
                 bottom_radius=(segment.bottomDiameter / 2),
                 top_radius=(segment.topDiameter / 2),
             )
+        case SquaredConeSegment():
+            return _volume_from_height_squared_cone(section_height, segment)
         case _:
             # TODO: implement volume calculations for truncated circular and rounded rectangular segments
             raise NotImplementedError(
@@ -275,6 +295,8 @@ def height_at_volume_within_section(
                 top_width=section.topXDimension,
                 top_length=section.topYDimension,
             )
+        case SquaredConeSegment():
+            return _height_from_volume_squared_cone(target_volume_relative, section)
         case _:
             raise NotImplementedError(
                 "Height from volume calculation not yet implemented for this well shape."
@@ -309,6 +331,8 @@ def volume_at_height_within_section(
                 top_width=section.topXDimension,
                 top_length=section.topYDimension,
             )
+        case SquaredConeSegment():
+            return _volume_from_height_squared_cone(target_height_relative, section)
         case _:
             # TODO(cm): this would be the NEST-96 2uL wells referenced in EXEC-712
             # we need to input the math attached to that issue

shared-data/python/Config.in

@@ -3,6 +3,7 @@ config BR2_PACKAGE_PYTHON_OPENTRONS_SHARED_DATA
   depends on BR2_PACKAGE_PYTHON3
   select BR2_PACKAGE_PYTHON_JSONSCHEMA # runtime
   select BR2_PACKAGE_PYTHON_TYPING_EXTENSIONS # runtime
+  select BR2_PACKAGE_PYTHON_NUMPY # runtime
 
   help
     Opentrons data sources. Used on an OT-2 robot.

shared-data/python/Pipfile

@@ -28,3 +28,4 @@ pytest-clarity = "~=1.0.0"
 opentrons-shared-data = { editable = true, path = "." }
 jsonschema = "==4.21.1"
 pydantic = "==1.10.12"
+numpy = "==1.22.3"

shared-data/python/opentrons_shared_data/labware/constants.py

@@ -8,9 +8,11 @@
 WELL_NAME_PATTERN: Final["re.Pattern[str]"] = re.compile(r"^([A-Z]+)([0-9]+)$", re.X)
 
 # These shapes are for wellshape definitions and describe the top of the well
-Circular = Literal["circular"]
-Rectangular = Literal["rectangular"]
-WellShape = Union[Circular, Rectangular]
+CircularType = Literal["circular"]
+Circular: CircularType = "circular"
+RectangularType = Literal["rectangular"]
+Rectangular: RectangularType = "rectangular"
+WellShape = Union[Literal["circular"], Literal["rectangular"]]
 
 # These shapes are used to describe the 3D primatives used to build wells
 Conical = Literal["conical"]

shared-data/python/opentrons_shared_data/labware/labware_definition.py

@@ -7,6 +7,9 @@
 
 from enum import Enum
 from typing import TYPE_CHECKING, Dict, List, Optional, Union
+from math import sqrt, asin
+from numpy import pi, trapz
+from functools import cached_property
 
 from pydantic import (
     BaseModel,
@@ -26,6 +29,8 @@
     SquaredCone,
     Spherical,
     WellShape,
+    Circular,
+    Rectangular,
 )
 
 SAFE_STRING_REGEX = "^[a-z0-9._]+$"
@@ -350,6 +355,87 @@ class SquaredConeSegment(BaseModel):
         description="The height at the bottom of a bounded subsection of a well, relative to the bottom of the well",
     )
 
+    @staticmethod
+    def _area_trap_points(
+        total_frustum_height: float,
+        circle_diameter: float,
+        rectangle_x: float,
+        rectangle_y: float,
+        dx: float,
+    ) -> List[float]:
+        """Grab a bunch of data points of area at given heights."""
+
+        def _area_arcs(r: float, c: float, d: float) -> float:
+            """Return the area of all 4 arc segments."""
+            theata_y = asin(c / r)
+            theata_x = asin(d / r)
+            theata_arc = (pi / 2) - theata_y - theata_x
+            # area of all 4 arcs is 4 * pi*r^2*(theata/2pi)
+            return 2 * r**2 * theata_arc
+
+        def _area(r: float) -> float:
+            """Return the area of a given r_y."""
+            # distance from the center of y axis of the rectangle to where the arc intercepts that side
+            c: float = (
+                sqrt(r**2 - (rectangle_y / 2) ** 2) if (rectangle_y / 2) < r else 0
+            )
+            # distance from the center of x axis of the rectangle to where the arc intercepts that side
+            d: float = (
+                sqrt(r**2 - (rectangle_x / 2) ** 2) if (rectangle_x / 2) < r else 0
+            )
+            arc_area = _area_arcs(r, c, d)
+            y_triangles: float = rectangle_y * c
+            x_triangles: float = rectangle_x * d
+            return arc_area + y_triangles + x_triangles
+
+        r_0 = circle_diameter / 2
+        r_h = sqrt(rectangle_x**2 + rectangle_y**2) / 2
+
+        num_steps = int(total_frustum_height / dx)
+        points = [0.0]
+        for i in range(num_steps + 1):
+            r_y = (i * dx / total_frustum_height) * (r_h - r_0) + r_0
+            points.append(_area(r_y))
+        return points
+
+    @cached_property
+    def height_to_volume_table(self) -> Dict[float, float]:
+        """Return a lookup table of heights to volumes."""
+        # the accuracy of this method is approximately +- 10*dx so for dx of 0.001 we have a +- 0.01 ul
+        dx = 0.001
+        total_height = self.topHeight - self.bottomHeight
+        points = SquaredConeSegment._area_trap_points(
+            total_height,
+            self.circleDiameter,
+            self.rectangleXDimension,
+            self.rectangleYDimension,
+            dx,
+        )
+        if self.bottomCrossSection is Rectangular:
+            # The points function assumes the circle is at the bottom but if its flipped we just reverse the points
+            points.reverse()
+        elif self.bottomCrossSection is not Circular:
+            raise NotImplementedError(
+                "If you see this error a new well shape has been added without updating this code"
+            )
+        y = 0.0
+        table: Dict[float, float] = {}
+        # fill in the table
+        while y < total_height:
+            table[y] = trapz(points[0 : int(y / dx)], dx=dx)
+            y = y + dx
+
+        # we always want to include the volume at the max height
+        table[total_height] = trapz(points, dx=dx)
+        return table
+
+    @cached_property
+    def volume_to_height_table(self) -> Dict[float, float]:
+        return dict((v, k) for k, v in self.height_to_volume_table.items())
+
+    class Config:
+        keep_untouched = (cached_property,)
+
 
 """
 module filitedCuboidSquare(bottom_shape, diameter, width, length, height, steps) {

shared-data/python/opentrons_shared_data/labware/types.py

@@ -7,8 +7,8 @@
 from typing_extensions import Literal, TypedDict, NotRequired
 from .labware_definition import InnerWellGeometry
 from .constants import (
-    Circular,
-    Rectangular,
+    CircularType,
+    RectangularType,
 )
 
 LabwareUri = NewType("LabwareUri", str)
@@ -84,7 +84,7 @@ class LabwareDimensions(TypedDict):
 
 
 class CircularWellDefinition(TypedDict):
-    shape: Circular
+    shape: CircularType
     depth: float
     totalLiquidVolume: float
     x: float
@@ -95,7 +95,7 @@ class CircularWellDefinition(TypedDict):
 
 
 class RectangularWellDefinition(TypedDict):
-    shape: Rectangular
+    shape: RectangularType
     depth: float
     totalLiquidVolume: float
     x: float