Fixing lots of ugly stuff

arm61 · Nov 23, 2022 · ed87127 · ed87127
1 parent e61eb97
commit ed87127
Show file tree

Hide file tree

Showing 2 changed files with 52 additions and 3 deletions.
diff --git a/uravu/distribution.py b/uravu/distribution.py
@@ -166,6 +166,15 @@ def con_int(self, ci_points: List[float]=[2.5, 97.5]) -> np.ndarray:
         """
         Get the extrema of the confidence intervals of the distribution.
 
+        :param ci_points: The confidence interval points to return.
+        :return: Distribution values at the confidence interval.
+        """
+        return self.ci(ci_points)
+
+    def ci(self, ci_points: List[float]=[2.5, 97.5]) -> np.ndarray:
+        """
+        Get the extrema of the confidence intervals of the distribution.
+
         :param ci_points: The confidence interval points to return.
         :return: Distribution values at the confidence interval.
         """
@@ -187,10 +196,50 @@ def __repr__(self) -> np.ndarray:
         """
         :return: Representation.
         """
+        return self.samples.__repr__()
+
+    def __array__(self) -> np.ndarray:
+        """
+        :return: Array of samples.
+        """
         return self.samples
 
+    def __add__(self, obj2: Union[float, np.ndarray]) -> np.ndarray:
+        """
+        :return: Sum of samples and obj2.
+        """
+        return np.add(self, obj2)
+
+    def __sub__(self, obj2: Union[float, np.ndarray]) -> np.ndarray:
+        """
+        :return: Difference between samples and obj2.
+        """
+        return np.subtract(self, obj2)
+
+    def __mul__(self, obj2: Union[float, np.ndarray]) -> np.ndarray:
+        """
+        :return: Product of samples and obj2.
+        """
+        return np.multiply(self, obj2)
+
+    def __truediv__(self, obj2: Union[float, np.ndarray]) -> np.ndarray:
+        """
+        :return: Ratio between samples and obj2.
+        """
+        return np.divide(self, obj2)
+
+    def __mod__(self, obj2: Union[float, np.ndarray]) -> np.ndarray:
+        """
+        :return: Modulus of samples by obj2.
+        """
+        return np.remainder(self, obj2)
+
     def __str__(self) -> np.ndarray:
         """
         :return: String representation.
         """
-        return self.samples
+        if self.s is not None:
+            return f'{self.n:.3e} +/- {self.s:.3e}'
+        else:
+            c = np.abs(self.n - self.con_int())
+            return f'{self.n:.3e}(+{c[1]:.3e}/-{c[0]:.3e}'
diff --git a/uravu/tests/test_distribution.py b/uravu/tests/test_distribution.py
@@ -99,11 +99,11 @@ def test_add_samples(self):
 
     def test_repr(self):
         distro = Distribution(norm.rvs(loc=0, scale=1, size=100, random_state=np.random.RandomState(2)))
-        assert_equal(distro.__repr__(), distro.samples)
+        assert_equal(distro.__repr__(), distro.samples.__repr__())
 
     def test_str(self):
         distro = Distribution(norm.rvs(loc=2, scale=0.0001, size=100, random_state=np.random.RandomState(2)))
-        assert_equal(distro.__str__(), distro.samples)
+        assert_equal(distro.__str__(), f'{distro.n:.3e} +/- {distro.s:.3e}')
 
     def test_dict_roundtrip(self):
         distro = Distribution(norm.rvs(loc=0, scale=1, size=1000), name='random')