fix and augment test_vax_beta

afmagee42 · afmagee42 · commit ef9a0a64cb0c · 2024-12-04T09:00:46.000-08:00
diff --git a/tests/test_ngm.py b/tests/test_ngm.py
@@ -28,12 +28,13 @@ def test_vax_beta():
     current = ngm.get_R(beta=beta, n=n, n_vax=n_vax, ve=ve)
     expected = np.array(
         [
-            [10.0 * .2 * .5, 0.1 * .8],
-            [0.1  * .2 * .5, 1.0 * .8],
+            [10.0 * .2 * .5, 0.1  * .2 * .5],
+            [0.1 * .8, 1.0 * .8],
         ]
     )
 
     assert_array_equal(current, expected)
+    assert ngm.dominant_eigen(current).value < 2.0013
 
 
 def test_simulate():
@@ -87,7 +88,7 @@ def test_kr():
         ve = 1.0,
     )
 
-    r_p_61 = np.array([[2, 0.02],[0.08, 0.8]]),
+    r_p_61 = np.array([[2, 0.02],[0.08, 0.8]])
     assert np.isclose(r, r_p_61).all()
 
     r0 = ngm.dominant_eigen(r).value

Original file line number	Diff line number	Diff line change
`@@ -28,12 +28,13 @@ def test_vax_beta():`
`28`	`28`	`current = ngm.get_R(beta=beta, n=n, n_vax=n_vax, ve=ve)`
`29`	`29`	`expected = np.array(`
`30`	`30`	`[`
`31`		`- [10.0 * .2 * .5, 0.1 * .8],`
`32`		`- [0.1 * .2 * .5, 1.0 * .8],`
	`31`	`+ [10.0 * .2 * .5, 0.1 * .2 * .5],`
	`32`	`+ [0.1 * .8, 1.0 * .8],`
`33`	`33`	`]`
`34`	`34`	`)`
`35`	`35`
`36`	`36`	`assert_array_equal(current, expected)`
	`37`	`+ assert ngm.dominant_eigen(current).value < 2.0013`
`37`	`38`
`38`	`39`
`39`	`40`	`def test_simulate():`
`@@ -87,7 +88,7 @@ def test_kr():`
`87`	`88`	`ve = 1.0,`
`88`	`89`	`)`
`89`	`90`
`90`		`- r_p_61 = np.array([[2, 0.02],[0.08, 0.8]]),`
	`91`	`+ r_p_61 = np.array([[2, 0.02],[0.08, 0.8]])`
`91`	`92`	`assert np.isclose(r, r_p_61).all()`
`92`	`93`
`93`	`94`	`r0 = ngm.dominant_eigen(r).value`