Added tests for table

xsuite · Mar 28, 2023 · 035a1bc · 035a1bc
1 parent 284316b
commit 035a1bc
Showing 1 changed file with 20 additions and 0 deletions.
diff --git a/tests/test_table.py b/tests/test_table.py
@@ -19,13 +19,16 @@ def test_column_selection():
     assert t["betx bety"].betx[0] == t.betx[0]
     assert t[["betx", "bety"]].betx[0] == t.betx[0]
     assert t["sqrt(betx)>3 sqrt(bety)"]['sqrt(betx)>3'][0]==(np.sqrt(data['betx'][0])>3)
+    assert isinstance(t.cols['betx'], RDMTable)
+    assert t.cols["betx", "bety"].betx[0] == t.betx[0]
 
 
 def test_row_selection():
     assert t[:, 1].betx==data['betx'][1]
     assert t[:, [0, 2]].betx[1]==data['betx'][2]
     assert t[:, t.s > 1].betx[1]==data['betx'][t.s > 1][1]
     assert t[:, t.s > 1, 1].betx[0]==data['betx'][t.s > 1][0]
+    assert t.rows[t.s > 1, 1].betx[0]==data['betx'][t.s > 1][0]
 
 def test_row_selection_names():
     assert t[:, "ip1"].betx[0]==data['betx'][0]
@@ -34,6 +37,13 @@ def test_row_selection_names():
     assert t[:, "notthere"]._nrows==0
     assert t[:, ["ip1", "ip2"]].betx[0]==data['betx'][0]
 
+def test_row_selection_names_with_rows():
+    assert t.rows["ip2"].betx[0]==data['betx'][1]
+    assert t.rows["ip[23]"].betx[0]==data['betx'][1]
+    assert t.rows["ip.*##1"].betx[0]==data['betx'][1]
+    assert t.rows["notthere"]._nrows==0
+    assert t.rows[["ip1", "ip2"]].betx[1]==data['betx'][1]
+
 def test_row_selection_ranges():
     assert t[:, 1:4:3].betx[0] == data['betx'][1]
     assert t[:, 1.5:2.5:"s"].betx[0] == data['betx'][1]
@@ -43,3 +53,13 @@ def test_row_selection_ranges():
     assert t[:, "ip1":"ip3":"name"].betx[0]  == data['betx'][0]
     assert t[:, None].betx[0]  == data['betx'][0]
     assert t[:, :].betx[0]  == data['betx'][0]
+
+def test_row_selection_ranges_with_rows():
+    assert t.rows[1:4:3].betx[0] == data['betx'][1]
+    assert t.rows[1.5:2.5:"s"].betx[0] == data['betx'][1]
+    assert t.rows["ip1":"ip3"].betx[2] == data['betx'][2]
+    assert t.rows["ip.*##1":"ip.*##2"].betx[0] == data['betx'][1]
+    assert t.rows["ip2%%-1":"ip2%%+1"].betx[0] == data['betx'][0]
+    assert t.rows["ip1":"ip3":"name"].betx[0]  == data['betx'][0]
+    assert t.rows[None].betx[0]  == data['betx'][0]
+    assert t.rows[:].betx[0]  == data['betx'][0]