Add tests for SAM matrix

open_cge/execute.py

@@ -23,6 +23,49 @@
 w = ('LAB', 'LAND', 'NTR')
 
 
+def check_square():
+    '''
+    this function tests whether the SAM is a square matrix.
+    '''
+    sam_small = sam.iloc[:, :-3]
+    sam_small = sam_small.drop("TOTAL")
+    sam_small.to_numpy(dtype=None, copy=True)
+    if not sam_small.shape[0] == sam_small.shape[1]:
+        raise ValueError(f"SAM is not square. It has {sam_small.shape[0]} rows and {sam_small.shape[0]} columns")
+
+def row_total():
+    '''
+    this function tests whether the row sums
+    of the SAM equal the expected value.
+    '''
+    sam_small = sam.iloc[:, :-3]
+    sam_small = sam_small.drop("TOTAL")
+    row_sum = sam_small.sum(axis=0)
+    row_sum = pd.Series(row_sum)
+    return row_sum
+
+def col_total():
+    '''
+    this function tests whether column sums
+    of the SAM equal the expected values.
+    '''
+    sam_small = sam.iloc[:, :-3]
+    sam_small = sam_small.drop("TOTAL")
+    col_sum = sam_small.sum(axis=1)
+    col_sum = pd.Series(col_sum)
+    return col_sum
+
+def row_col_equal():
+    '''
+    this function tests whether row sums
+    and column sums of the SAM are equal.
+    '''
+    sam_small = sam.iloc[:, :-3]
+    sam_small = sam_small.drop("TOTAL")
+    row_sum = sam_small.sum(axis=0)
+    col_sum = sam_small.sum(axis=1)
+    np.testing.assert_allclose(row_sum, col_sum)
+
 def runner():
     '''
     this function runs the cge model
@@ -107,4 +150,6 @@ def runner():
 
 
 if __name__ == "__main__":
+    check_square()
+    row_col_equal()
     runner()

open_cge/tests/test_execute.py

@@ -6,11 +6,59 @@
 
 def test_runner():
     '''
-    Test of mdoel solution from execute.runner
+    Test of model solution from execute.runner
     '''
 
     dict = {'AGR': 115.999989, 'OIL': 149.000054, 'IND': 1431.999758, 'SER': 1803.000214}
     expectedQ = pd.DataFrame.from_dict(dict, orient='index')
     expectedQ.columns = [None]
     testQ = execute.runner()
     assert_series_equal(expectedQ[None], testQ, check_dtype=False)
+
+def test_check_square():
+    '''
+    Test of check_square function
+    '''
+    testdata = [[1, 2, 3], [2, 3, 4], [3, 1, 4]]
+    mat = pd.DataFrame(data = testdata, index = None)
+    mat.to_numpy(dtype=None, copy=True)
+    if not mat.shape[0] == mat.shape[1]:
+        raise ValueError(f"SAM is not square. It has {mat.shape[0]} rows and {mat.shape[1]} columns")
+
+
+def test_row_total():
+    '''
+    Test of row_total function
+    '''
+    rowdict = {'AGR': 118, 'OIL': 1015.435052, 'IND': 1738, 'SER': 1837, 'LAB': 551,
+    'CAP': 1051, 'LAND': 24, 'NTR': 890.4350515, 'DTX': 948, 'IDT': 21, 'ACT': 9,
+    'HOH': 2616.435052, 'GOV': 978, 'INV': 485, 'EXT': 1504.435052}
+    expected_rowtotal =  pd.DataFrame.from_dict(rowdict, orient='index')
+    expected_rowtotal.columns = [None]
+    test_rowtotal = execute.row_total()
+    assert_series_equal(expected_rowtotal[None], test_rowtotal, check_dtype=False)
+
+
+def test_col_total():
+    '''
+    Test of col_total function
+    '''
+    coldict = {'AGR': 118, 'OIL': 1015.435052, 'IND': 1738, 'SER': 1837, 'LAB': 551,
+    'CAP': 1051, 'LAND': 24, 'NTR': 890.4350515, 'DTX': 948, 'IDT': 21, 'ACT': 9,
+    'HOH': 2616.435052, 'GOV': 978, 'INV': 485, 'EXT': 1504.435052}
+    expected_coltotal =  pd.DataFrame.from_dict(coldict, orient='index')
+    expected_coltotal.columns = [None]
+    test_coltotal = execute.col_total()
+    assert_series_equal(expected_coltotal[None], test_coltotal, check_dtype=False)
+
+
+def test_row_col_equal():
+    '''
+    Test of row_col_equal function
+    '''
+    data = {'row_1': [3, 2, 1, 0], 'row_2': [2, 3, 4, 5], 'row_3': [1, 4, 3, 6],
+    'row_4': [0, 5, 6, 3]}
+    test_df = pd.DataFrame.from_dict(data, orient='index')
+    test_row_sum = test_df.sum(axis=0)
+    test_col_sum = test_df.sum(axis=1)
+    np.testing.assert_allclose(test_row_sum, test_col_sum)