diff --git a/src/SymptomsCausedByVaccines/DataFrameFilter.py b/src/SymptomsCausedByVaccines/DataFrameFilter.py index fca5456a713..141c5449eeb 100644 --- a/src/SymptomsCausedByVaccines/DataFrameFilter.py +++ b/src/SymptomsCausedByVaccines/DataFrameFilter.py @@ -1,9 +1,17 @@ class DataFrameFilter: @staticmethod - def withoutZeroRows(dataFrame): + def withoutZeroRowsAndZeroColumns(dataFrame): + return DataFrameFilter._withoutZeroColumns(DataFrameFilter._withoutZeroRows(dataFrame)) + + @staticmethod + def _withoutZeroRows(dataFrame): return dataFrame.loc[~DataFrameFilter._isZeroRow(dataFrame)] @staticmethod def _isZeroRow(dataFrame): return (dataFrame == 0.0).all(axis = 'columns') + + @staticmethod + def _withoutZeroColumns(dataFrame): + return dataFrame.loc[:, (dataFrame != 0.0).any(axis = 'index')] diff --git a/src/SymptomsCausedByVaccines/DataFrameFilterTest.py b/src/SymptomsCausedByVaccines/DataFrameFilterTest.py index 687f557cb30..67efc3c6e28 100644 --- a/src/SymptomsCausedByVaccines/DataFrameFilterTest.py +++ b/src/SymptomsCausedByVaccines/DataFrameFilterTest.py @@ -6,6 +6,34 @@ from SymptomsCausedByVaccines.DataFrameFilter import DataFrameFilter class DataFrameFilterTest(unittest.TestCase): + def test_withoutZeroRowsAndZeroColumns(self): + # Given + dataFrame = TestHelper.createDataFrame( + columns = ['col1', 'zero column'], + data = [ [0.6, 0.0], + [0.0, 0.0]], + index = pd.Index( + name = 'VAX_TYPE', + data = [ + '6VAX-F', + 'ZERO ROW' + ])) + + # When + dataFrameWithoutZeroRowsAndZeroColumns = DataFrameFilter.withoutZeroRowsAndZeroColumns(dataFrame) + + # Then + assert_frame_equal( + dataFrameWithoutZeroRowsAndZeroColumns, + TestHelper.createDataFrame( + columns = ['col1'], + data = [ [0.6]], + index = pd.Index( + name = 'VAX_TYPE', + data = [ + '6VAX-F' + ]))) + def test_withoutZeroRows_noZeroRow(self): # Given dataFrame = TestHelper.createDataFrame( @@ -20,7 +48,7 @@ class DataFrameFilterTest(unittest.TestCase): ])) # When - dataFrameWithoutZeroRows = DataFrameFilter.withoutZeroRows(dataFrame) + dataFrameWithoutZeroRows = DataFrameFilter._withoutZeroRows(dataFrame) # Then assert_frame_equal(dataFrameWithoutZeroRows, dataFrame) @@ -39,7 +67,7 @@ class DataFrameFilterTest(unittest.TestCase): ])) # When - dataFrameWithoutZeroRows = DataFrameFilter.withoutZeroRows(dataFrame) + dataFrameWithoutZeroRows = DataFrameFilter._withoutZeroRows(dataFrame) # Then assert_frame_equal( @@ -52,3 +80,52 @@ class DataFrameFilterTest(unittest.TestCase): data = [ '6VAX-F' ]))) + + def test_withoutZeroColumns_noZeroColumn(self): + # Given + dataFrame = TestHelper.createDataFrame( + columns = ['col1', 'col2'], + data = [ [0.6, 1.5], + [0.3, 3.0]], + index = pd.Index( + name = 'VAX_TYPE', + data = [ + '6VAX-F', + 'ADEN' + ])) + + # When + dataFrameWithoutZeroColumns = DataFrameFilter._withoutZeroColumns(dataFrame) + + # Then + assert_frame_equal(dataFrameWithoutZeroColumns, dataFrame) + + def test_withoutZeroColumns(self): + # Given + dataFrame = TestHelper.createDataFrame( + columns = ['col1', 'zero column'], + data = [ [0.6, 0.0], + [1.2, 0.0]], + index = pd.Index( + name = 'VAX_TYPE', + data = [ + '6VAX-F', + '6VAX-G' + ])) + + # When + dataFrameWithoutZeroColumns = DataFrameFilter._withoutZeroColumns(dataFrame) + + # Then + assert_frame_equal( + dataFrameWithoutZeroColumns, + TestHelper.createDataFrame( + columns = ['col1'], + data = [ [0.6], + [1.2]], + index = pd.Index( + name = 'VAX_TYPE', + data = [ + '6VAX-F', + '6VAX-G' + ])))