Repository: manmartgarc/stochatreat

Issue:
Rounding error in assignment probabilities
I noticed that the check for whether the assignment probabilities passed to the `probs` argument is equal to 1 can fail due to a rounding error. As expected, this happens because fractions are converted to floats, which introduces some imprecision. The code excerpt below is where the issue lies:

https://github.com/manmartgarc/stochatreat/blob/38d0e8d087c11178a6dd5ecd7d518527153bf945/src/stochatreat/stochatreat.py#L99-L102

Consider `probs = [1/2, 1/3, 1/6]`. These fractions add up to 1. However, `np.array(probs)` yields `array([0.5       , 0.33333333, 0.16666667])`. Therefore, `np.array(probs).sum()` yields `0.9999999999999999 != 1`.

There are several ways to avoid this issue. For instance, we can use the `Fraction` class from the [`fractions`](https://docs.python.org/3/library/fractions.html) built-in library:

```python
import numpy as np
from fractions import Fraction

probs = [1/2, 1/3, 1/6]
probs_np = np.array(probs)
probs_sum = float(np.array([Fraction(f).limit_denominator() for f in probs]).sum())

probs_sum != 1
>>> False
```

Discussion:

Patch:
diff --git a/src/stochatreat/__about__.py b/src/stochatreat/__about__.py
index 39e62e7..5d7942c 100644
--- a/src/stochatreat/__about__.py
+++ b/src/stochatreat/__about__.py
@@ -1,2 +1,2 @@
 # pragma: no cover
-__version__ = "0.0.19"
+__version__ = "0.0.20"
diff --git a/src/stochatreat/stochatreat.py b/src/stochatreat/stochatreat.py
index 58949dc..a913a5a 100644
--- a/src/stochatreat/stochatreat.py
+++ b/src/stochatreat/stochatreat.py
@@ -1,5 +1,15 @@
+"""Stratified random assignment of treatments to units.
+
+This module provides a function to assign treatments to units in a
+stratified manner. The function is designed to work with pandas
+dataframes and is able to handle multiple strata. There are also different
+strategies to deal with misfits (units that are left over after the
+stratified assignment procedure).
+"""
+
 from __future__ import annotations
 
+import math
 from typing import Literal
 
 import numpy as np
@@ -90,7 +100,7 @@ def stochatreat(
         probs_np = np.array([frac] * len(treatment_ids))
     elif probs is not None:
         probs_np = np.array(probs)
-        if probs_np.sum() != 1:
+        if not math.isclose(probs_np.sum(), 1, rel_tol=1e-9):
             error_msg = "The probabilities must add up to 1"
             raise ValueError(error_msg)
         if len(probs_np) != len(treatment_ids):


Test patch:
diff --git a/tests/test_assignment.py b/tests/test_assignment.py
index e079b2f..f8f2749 100644
--- a/tests/test_assignment.py
+++ b/tests/test_assignment.py
@@ -30,6 +30,7 @@ standard_probs = [
     [0.5, 0.5],
     [2 / 3, 1 / 3],
     [0.9, 0.1],
+    [1 / 2, 1 / 3, 1 / 6],
 ]
 
 # a set of stratum column combinations from the above df fixture to throw at