Ejemplo n.º 1
0
 def test_where(self):
     helpers.assert_quantity_equal(
         np.where(self.q >= 2 * self.ureg.m, self.q, 20 * self.ureg.m),
         [[20, 2], [3, 4]] * self.ureg.m,
     )
     helpers.assert_quantity_equal(
         np.where(self.q >= 2 * self.ureg.m, self.q, 0),
         [[0, 2], [3, 4]] * self.ureg.m,
     )
     helpers.assert_quantity_equal(
         np.where(self.q >= 2 * self.ureg.m, self.q, np.nan),
         [[np.nan, 2], [3, 4]] * self.ureg.m,
     )
     helpers.assert_quantity_equal(
         np.where(self.q >= 3 * self.ureg.m, 0, self.q),
         [[1, 2], [0, 0]] * self.ureg.m,
     )
     helpers.assert_quantity_equal(
         np.where(self.q >= 3 * self.ureg.m, np.nan, self.q),
         [[1, 2], [np.nan, np.nan]] * self.ureg.m,
     )
     helpers.assert_quantity_equal(
         np.where(self.q >= 2 * self.ureg.m, self.q, np.array(np.nan)),
         [[np.nan, 2], [3, 4]] * self.ureg.m,
     )
     helpers.assert_quantity_equal(
         np.where(self.q >= 3 * self.ureg.m, np.array(np.nan), self.q),
         [[1, 2], [np.nan, np.nan]] * self.ureg.m,
     )
     with pytest.raises(DimensionalityError):
         np.where(
             self.q < 2 * self.ureg.m,
             self.q,
             0 * self.ureg.J,
         )
Ejemplo n.º 2
0
 def test_where(self):
     self.assertQuantityEqual(
         np.where(self.q >= 2 * self.ureg.m, self.q, 20 * self.ureg.m),
         [[20, 2], [3, 4]] * self.ureg.m,
     )
     self.assertQuantityEqual(
         np.where(self.q >= 2 * self.ureg.m, self.q, 0),
         [[0, 2], [3, 4]] * self.ureg.m,
     )
     self.assertQuantityEqual(
         np.where(self.q >= 2 * self.ureg.m, self.q, np.nan),
         [[np.nan, 2], [3, 4]] * self.ureg.m,
     )
     self.assertQuantityEqual(
         np.where(self.q >= 3 * self.ureg.m, 0, self.q),
         [[1, 2], [0, 0]] * self.ureg.m,
     )
     self.assertQuantityEqual(
         np.where(self.q >= 3 * self.ureg.m, np.nan, self.q),
         [[1, 2], [np.nan, np.nan]] * self.ureg.m,
     )
     self.assertQuantityEqual(
         np.where(self.q >= 2 * self.ureg.m, self.q, np.array(np.nan)),
         [[np.nan, 2], [3, 4]] * self.ureg.m,
     )
     self.assertQuantityEqual(
         np.where(self.q >= 3 * self.ureg.m, np.array(np.nan), self.q),
         [[1, 2], [np.nan, np.nan]] * self.ureg.m,
     )
     self.assertRaises(
         DimensionalityError,
         np.where,
         self.q < 2 * self.ureg.m,
         self.q,
         0 * self.ureg.J,
     )