Unit Test for aggregate_spot_compositions

[1]:
import numpy as np
import pandas as pd
import unittest
from pandas.testing import assert_frame_equal

import os
os.sys.path.append('../../../')
from mesa.ecospatial._utils import aggregate_spot_compositions
[9]:
class TestAggregateSpotCompositions(unittest.TestCase):
    def test_basic(self):
        """
        Test a basic scenario with a small 2x2 labelled array.
        """
        # labelled: 2x2 array with:
        #   - (0,0)=0 (background), (0,1)=1, (1,0)=2, (1,1)=2
        labelled = np.array([[0, 1],
                             [2, 2]])

        compositions = [
            pd.Series([1.0, 2.0], index=['A', 'B']),  # position (0,0) -> background, ignored
            pd.Series([3.0, 4.0], index=['A', 'B']),  # position (0,1) -> Island 1
            pd.Series([5.0, 6.0], index=['A', 'B']),  # position (1,0) -> Island 2
            pd.Series([7.0, 8.0], index=['A', 'B'])   # position (1,1) -> Island 2
        ]

        # Expected result:
        # - Island_1: composition = [3, 4]
        # - Island_2: composition = [5+7, 6+8] = [12, 14]
        expected_data = {
            'Island_1': pd.Series([3.0, 4.0], index=['A', 'B']),
            'Island_2': pd.Series([12.0, 14.0], index=['A', 'B'])
        }
        expected_df = pd.DataFrame.from_dict(expected_data, orient='index')

        result = aggregate_spot_compositions(labelled, compositions)
        # Sort the columns for a consistent order before comparing
        result = result.sort_index(axis=1)
        expected_df = expected_df.sort_index(axis=1)
        assert_frame_equal(result, expected_df)

    def test_with_none(self):
        """
        Test with some spots having None for composition.
        """
        # labelled: 3x3 array with islands 1 and 2
        labelled = np.array([
            [0, 1, 1],
            [2, 2, 0],
            [1, 2, 2]
        ])

        # There are 9 spots (flattened indices 0-8). We assign:
        # - Spots with label 0: indices 0 and 5 -> None.
        # - Spots with label 1: indices 1, 2, 6.
        # - Spots with label 2: indices 3, 4, 7, 8.
        comps = []
        # index 0: (0,0) label 0 -> None
        comps.append(None)
        # index 1: (0,1) label 1 -> pd.Series({'X': 1})
        comps.append(pd.Series({'X': 1.0}))
        # index 2: (0,2) label 1 -> pd.Series({'X': 2, 'Y': 3})
        comps.append(pd.Series({'Y': 3.0}))
        # index 3: (1,0) label 2 -> pd.Series({'Y': 4})
        comps.append(pd.Series({'Y': 4.0}))
        # index 4: (1,1) label 2 -> None
        comps.append(None)
        # index 5: (1,2) label 0 -> None
        comps.append(None)
        # index 6: (2,0) label 1 -> pd.Series({'Z': 5})
        comps.append(pd.Series({'Z': 5.0}))
        # index 7: (2,1) label 2 -> pd.Series({'X': 3, 'Z': 2})
        comps.append(pd.Series({'X': 3.0, 'Z': 2.0}))
        # index 8: (2,2) label 2 -> pd.Series({'Y': 1})
        comps.append(pd.Series({'Y': 1.0}))

        # Expected aggregation:
        # Island_1 (indices 1, 2, 6):
        #   Sum: {'X': 1} + {'X': 2, 'Y': 3} + {'Z': 5} = {'X': 3, 'Y': 3, 'Z': 5}
        # Island_2 (indices 3, 4, 7, 8):
        #   Sum: {'Y': 4} + 0 + {'X': 3, 'Z': 2} + {'Y': 1} = {'X': 3, 'Y': 5, 'Z': 2}
        expected_data = {
            'Island_1': pd.Series({'X': 1.0, 'Y': 3.0, 'Z': 5.0}),
            'Island_2': pd.Series({'X': 3.0, 'Y': 5.0, 'Z': 2.0})
        }
        expected_df = pd.DataFrame.from_dict(expected_data, orient='index')
        result = aggregate_spot_compositions(labelled, comps)
        print(result)
        # Sort columns for consistent ordering
        result = result.sort_index(axis=1)
        expected_df = expected_df.sort_index(axis=1)
        assert_frame_equal(result, expected_df)

    def test_all_background(self):
        """
        Test the case where all spots are background (label 0).
        """
        labelled = np.zeros((2, 2), dtype=int)
        compositions = [pd.Series({'A': 1}) for _ in range(4)]
        result = aggregate_spot_compositions(labelled, compositions)
        # Expect an empty DataFrame (no islands)
        self.assertEqual(result.shape[0], 0)

    def test_union_of_keys(self):
        """
        Test that the union of composition keys is maintained across islands,
        even if an island does not have a composition for some keys.
        """
        # Island 1 has only key 'A', while Island 2 has only key 'B'.
        labelled = np.array([[1, 2]])
        compositions = [
            pd.Series({'A': 5}),  # For Island_1 (only key 'A')
            pd.Series({'B': 7})   # For Island_2 (only key 'B')
        ]
        # The expected aggregated DataFrame should include both keys for each island:
        expected_data = {
            'Island_1': pd.Series({'A': 5, 'B': 0}),
            'Island_2': pd.Series({'A': 0, 'B': 7})
        }
        expected_df = pd.DataFrame.from_dict(expected_data, orient='index')
        expected_df = expected_df.sort_index(axis=1)

        result = aggregate_spot_compositions(labelled, compositions)
        result = result.sort_index(axis=1)
        assert_frame_equal(result, expected_df)

# Run the tests in Jupyter Notebook
if __name__ == '__main__':
    unittest.main(argv=[''], verbosity=2, exit=False)

test_all_background (__main__.TestAggregateSpotCompositions.test_all_background)
Test the case where all spots are background (label 0). ... ok
test_basic (__main__.TestAggregateSpotCompositions.test_basic)
Test a basic scenario with a small 2x2 labelled array. ... ok
test_union_of_keys (__main__.TestAggregateSpotCompositions.test_union_of_keys)
Test that the union of composition keys is maintained across islands, ... FAIL
test_with_none (__main__.TestAggregateSpotCompositions.test_with_none)
Test with some spots having None for composition. ... ok

======================================================================
FAIL: test_union_of_keys (__main__.TestAggregateSpotCompositions.test_union_of_keys)
Test that the union of composition keys is maintained across islands,
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/var/folders/7g/phdhh_ld3dlbnrst0t60bwzr0000gn/T/ipykernel_75522/4104004902.py", line 116, in test_union_of_keys
    assert_frame_equal(result, expected_df)
  File "/opt/miniconda3/envs/mesa/lib/python3.11/site-packages/pandas/_testing/asserters.py", line 1148, in assert_frame_equal
    raise_assert_detail(
  File "/opt/miniconda3/envs/mesa/lib/python3.11/site-packages/pandas/_testing/asserters.py", line 598, in raise_assert_detail
    raise AssertionError(msg)
AssertionError: DataFrame are different

DataFrame shape mismatch
[left]:  (2, 1)
[right]: (2, 2)

----------------------------------------------------------------------
Ran 4 tests in 0.097s

FAILED (failures=1)
            X    Y    Z
Island_1  1.0  3.0  5.0
Island_2  3.0  5.0  2.0
[ ]: