add more tests

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

@@ -17,14 +17,15 @@ def reject_unacceptable_heights(
     potential_heights: List[float], max_height: float
 ) -> float:
     """Reject any solutions to a polynomial equation that cannot be the height of a frustum."""
-    valid_heights = []
+    valid_heights_set = set()
     for root in potential_heights:
         # reject any heights that are negative or greater than the max height
         if not iscomplex(root):
             # take only the real component of the root and round to 4 decimal places
             rounded_root = round(real(root), 4)
             if (rounded_root <= max_height) and (rounded_root >= 0):
-                valid_heights.append(rounded_root)
+                valid_heights_set.add(rounded_root)
+    valid_heights = [height for height in valid_heights_set]
     if len(valid_heights) != 1:
         raise InvalidLiquidHeightFound(
             message="Unable to estimate valid liquid height from volume."

api/tests/opentrons/protocols/geometry/test_frustum_helpers.py

@@ -1,10 +1,8 @@
 import pytest
-from math import pi, sqrt
+from math import pi, sqrt, isclose
 from typing import Any, List
 
-# make some fake frusta
 from opentrons_shared_data.labware.types import (
-    WellDefinition,
     RectangularBoundedSection,
     CircularBoundedSection,
     SphericalSegment,
@@ -13,53 +11,23 @@
     cross_section_area_rectangular,
     cross_section_area_circular,
     reject_unacceptable_heights,
+    get_boundary_cross_sections,
+    get_cross_section_area,
+    volume_from_frustum_formula,
+    circular_frustum_polynomial_roots,
+    rectangular_frustum_polynomial_roots,
+    volume_from_height_rectangular,
+    volume_from_height_circular,
+    volume_from_height_spherical,
+    height_from_volume_circular,
+    height_from_volume_rectangular,
+    height_from_volume_spherical,
 )
 from opentrons.protocol_engine.errors.exceptions import InvalidLiquidHeightFound
 
 
-@pytest.mark.parametrize(
-    ["max_height", "potential_heights", "expected_heights"],
-    [
-        (34, [complex(4, 5), complex(5, 0), 35, 34, 33, 10, 0], [5, 34, 33, 10, 0]),
-        (2934, [complex(4, 5), complex(5, 0)], [5]),
-        (100, [-99, -1, complex(99.99, 0), 101], [99.99]),
-        (2, [0, -1, complex(-1.5, 0)], [0]),
-        (8, [complex(7, 1), -0.01], []),
-    ],
-)
-def test_reject_unacceptable_heights(
-    max_height: float, potential_heights: List[Any], expected_heights: List[float]
-) -> None:
-    if len(expected_heights) != 1:
-        with pytest.raises(InvalidLiquidHeightFound):
-            reject_unacceptable_heights(
-                max_height=max_height, potential_heights=potential_heights
-            )
-    else:
-        found_heights = reject_unacceptable_heights(
-            max_height=max_height, potential_heights=potential_heights
-        )
-        assert found_heights == expected_heights[0]
-
-
-@pytest.mark.parametrize("diameter", [2, 5, 8, 356, 1000])
-def test_cross_section_area_circular(diameter: float) -> None:
-    expected_area = pi * (diameter / 2) ** 2
-    assert cross_section_area_circular(diameter) == expected_area
-
-
-@pytest.mark.parametrize(
-    ["x_dimension", "y_dimension"], [(1, 38402), (234, 983), (94857, 40), (234, 999)]
-)
-def test_cross_section_area_rectangular(x_dimension: float, y_dimension: float) -> None:
-    expected_area = x_dimension * y_dimension
-    assert (
-        cross_section_area_rectangular(x_dimension=x_dimension, y_dimension=y_dimension)
-        == expected_area
-    )
-
-
 def fake_frusta() -> List[List[Any]]:
+    """A bunch of weird fake well shapes."""
     frusta = []
     frusta.append(
         [
@@ -110,25 +78,209 @@ def fake_frusta() -> List[List[Any]]:
     frusta.append(
         [
             RectangularBoundedSection(
-                shape="rectangular", xDimension=27.0, yDimension=36.0, topHeight=0.0
+                shape="rectangular", xDimension=27.0, yDimension=36.0, topHeight=3.5
             ),
             RectangularBoundedSection(
-                shape="rectangular", xDimension=36.0, yDimension=26.0, topHeight=0.0
+                shape="rectangular", xDimension=36.0, yDimension=26.0, topHeight=1.5
             ),
             SphericalSegment(shape="spherical", radiusOfCurvature=4.0, depth=1.5),
         ]
     )
     return frusta
 
 
-# cross_section_circle test
-# cross rectangle test
-# volume of a frustum formula test
-# polynomial roots test circular
-# polynomial roots test rectangular
-# volume from height circular
-# volume from height rectangular
-# volume from height spherical
-# height from volume circular
-# height from volume rectangular
-# height from volume spherical
+@pytest.mark.parametrize(
+    ["max_height", "potential_heights", "expected_heights"],
+    [
+        (34, [complex(4, 5), complex(5, 0), 35, 34, 33, 10, 0], [5, 34, 33, 10, 0]),
+        (2934, [complex(4, 5), complex(5, 0)], [5]),
+        (100, [-99, -1, complex(99.99, 0), 101], [99.99]),
+        (2, [0, -1, complex(-1.5, 0)], [0]),
+        (8, [complex(7, 1), -0.01], []),
+    ],
+)
+def test_reject_unacceptable_heights(
+    max_height: float, potential_heights: List[Any], expected_heights: List[float]
+) -> None:
+    """Make sure we reject all mathematical solutions that are physically not possible."""
+    if len(expected_heights) != 1:
+        with pytest.raises(InvalidLiquidHeightFound):
+            reject_unacceptable_heights(
+                max_height=max_height, potential_heights=potential_heights
+            )
+    else:
+        found_heights = reject_unacceptable_heights(
+            max_height=max_height, potential_heights=potential_heights
+        )
+        assert found_heights == expected_heights[0]
+
+
+@pytest.mark.parametrize("diameter", [2, 5, 8, 356, 1000])
+def test_cross_section_area_circular(diameter: float) -> None:
+    """Test circular area calculation."""
+    expected_area = pi * (diameter / 2) ** 2
+    assert cross_section_area_circular(diameter) == expected_area
+
+
+@pytest.mark.parametrize(
+    ["x_dimension", "y_dimension"], [(1, 38402), (234, 983), (94857, 40), (234, 999)]
+)
+def test_cross_section_area_rectangular(x_dimension: float, y_dimension: float) -> None:
+    """Test rectangular area calculation."""
+    expected_area = x_dimension * y_dimension
+    assert (
+        cross_section_area_rectangular(x_dimension=x_dimension, y_dimension=y_dimension)
+        == expected_area
+    )
+
+
+@pytest.mark.parametrize("well", fake_frusta())
+def test_get_cross_section_boundaries(well: List[List[Any]]) -> None:
+    """Make sure get_cross_section_boundaries returns the expected list indices."""
+    i = 0
+    for f, next_f in get_boundary_cross_sections(well):
+        assert f == well[i]
+        assert next_f == well[i + 1]
+        i += 1
+
+
+@pytest.mark.parametrize("well", fake_frusta())
+def test_frustum_formula_volume(well: List[Any]) -> None:
+    """Test volume-of-a-frustum formula calculation."""
+    for f, next_f in get_boundary_cross_sections(well):
+        if f["shape"] == "spherical" or next_f["shape"] == "spherical":
+            # not going to use formula on spherical segments
+            continue
+        f_area = get_cross_section_area(f)
+        next_f_area = get_cross_section_area(next_f)
+        frustum_height = next_f["topHeight"] - f["topHeight"]
+        expected_volume = (f_area + next_f_area + sqrt(f_area * next_f_area)) * (
+            frustum_height / 3
+        )
+        found_volume = volume_from_frustum_formula(
+            area_1=f_area, area_2=next_f_area, height=frustum_height
+        )
+        assert found_volume == expected_volume
+
+
+@pytest.mark.parametrize("well", fake_frusta())
+def test_volume_and_height_circular(well: List[Any]) -> None:
+    """Test both volume and height calculations for circular frusta."""
+    if well[-1]["shape"] == "spherical":
+        return
+    total_height = well[0]["topHeight"]
+    for f, next_f in get_boundary_cross_sections(well):
+        if f["shape"] == next_f["shape"] == "circular":
+            top_radius = next_f["diameter"] / 2
+            bottom_radius = f["diameter"] / 2
+            a = pi * ((top_radius - bottom_radius) ** 2) / (3 * total_height**2)
+            b = pi * bottom_radius * (top_radius - bottom_radius) / total_height
+            c = pi * bottom_radius**2
+            assert circular_frustum_polynomial_roots(
+                top_radius=top_radius,
+                bottom_radius=bottom_radius,
+                total_frustum_height=total_height,
+            ) == (a, b, c)
+            # test volume within a bunch of arbitrary heights
+            for target_height in range(round(total_height)):
+                expected_volume = (
+                    a * (target_height**3)
+                    + b * (target_height**2)
+                    + c * target_height
+                )
+                found_volume = volume_from_height_circular(
+                    target_height=target_height,
+                    total_frustum_height=total_height,
+                    bottom_radius=bottom_radius,
+                    top_radius=top_radius,
+                )
+                assert found_volume == expected_volume
+                # test going backwards to get height back
+                found_height = height_from_volume_circular(
+                    volume=found_volume,
+                    total_frustum_height=total_height,
+                    bottom_radius=bottom_radius,
+                    top_radius=top_radius,
+                )
+                assert isclose(found_height, target_height)
+
+
+@pytest.mark.parametrize("well", fake_frusta())
+def test_volume_and_height_rectangular(well: List[Any]) -> None:
+    """Test both volume and height calculations for rectangular frusta."""
+    if well[-1]["shape"] == "spherical":
+        return
+    total_height = well[0]["topHeight"]
+    for f, next_f in get_boundary_cross_sections(well):
+        if f["shape"] == next_f["shape"] == "rectangular":
+            top_length = next_f["yDimension"]
+            top_width = next_f["xDimension"]
+            bottom_length = f["yDimension"]
+            bottom_width = f["xDimension"]
+            a = (
+                (top_length - bottom_length)
+                * (top_width - bottom_width)
+                / (3 * total_height**2)
+            )
+            b = (
+                (bottom_length * (top_width - bottom_width))
+                + (bottom_width * (top_length - bottom_length))
+            ) / (2 * total_height)
+            c = bottom_length * bottom_width
+            assert rectangular_frustum_polynomial_roots(
+                top_length=top_length,
+                bottom_length=bottom_length,
+                top_width=top_width,
+                bottom_width=bottom_width,
+                total_frustum_height=total_height,
+            ) == (a, b, c)
+            # test volume within a bunch of arbitrary heights
+            for target_height in range(round(total_height)):
+                expected_volume = (
+                    a * (target_height**3)
+                    + b * (target_height**2)
+                    + c * target_height
+                )
+                found_volume = volume_from_height_rectangular(
+                    target_height=target_height,
+                    total_frustum_height=total_height,
+                    bottom_length=bottom_length,
+                    bottom_width=bottom_width,
+                    top_length=top_length,
+                    top_width=top_width,
+                )
+                assert found_volume == expected_volume
+                # test going backwards to get height back
+                found_height = height_from_volume_rectangular(
+                    volume=found_volume,
+                    total_frustum_height=total_height,
+                    bottom_length=bottom_length,
+                    bottom_width=bottom_width,
+                    top_length=top_length,
+                    top_width=top_width,
+                )
+                assert isclose(found_height, target_height)
+
+
+@pytest.mark.parametrize("well", fake_frusta())
+def test_volume_and_height_spherical(well: List[Any]) -> None:
+    """Test both volume and height calculations for spherical segments."""
+    if well[0]["shape"] == "spherical":
+        for target_height in range(round(well[0]["depth"])):
+            expected_volume = (
+                (1 / 3)
+                * pi
+                * (target_height**2)
+                * (3 * well[0]["radiusOfCurvature"] - target_height)
+            )
+            found_volume = volume_from_height_spherical(
+                target_height=target_height,
+                radius_of_curvature=well[0]["radiusOfCurvature"],
+            )
+            assert found_volume == expected_volume
+            found_height = height_from_volume_spherical(
+                volume=found_volume,
+                radius_of_curvature=well[0]["radiusOfCurvature"],
+                total_frustum_height=well[0]["depth"],
+            )
+            assert isclose(found_height, target_height)