Tidy up and docs

Author: danielward27

paramax/wrappers.py

@@ -161,29 +161,36 @@ def unwrap(self) -> T:
 
 
 class RealToIncreasingOnInterval(AbstractUnwrappable[Array]):
-    """Unconstrained vector to increasing on a fixed interval.
+    """Map an unconstrained vector to increasing points on a fixed interval.
 
-    Unconstrained vector is passed into softmax to obtain widths,
-    which are cumulatively summed and scaled to fit the interval
-    (or the remainder not used by min_width).
+    The input vector is transformed via a softmax into positive widths, to fill
+    the interval after adding minimum width. The cumulative sum of the widths
+    produces the incresing points.
 
-    Note an array of size d parameterizes widths, so maps to an array
-    of size d+1 if both endpoints are included, d if one is included,
-    and d-1 if neither are included.
+    If an array of size d is used, the result has size d+1 if both
+    endpoints are included, d if one is included, and d-1 if neither
+    are included.
+
+    Args:
+        arr: Unconstrained vector parameterizing the widths.
+        interval: (lower, upper) bounds of the interval.
+        min_width: Minimum spacing between consecutive points.
+        include_endpoints: Which endpoints to include: "both", "neither",
+            "lower", or "upper".
     """
 
     arr: Array
     interval: tuple[float | int, float | int]
     min_width: float
-    include_ends: Literal["both", "neither", "lower", "upper"]
+    include_endpoints: Literal["both", "neither", "lower", "upper"]
 
     def __init__(
         self,
         arr: Array,
         interval: tuple[float | int, float | int],
         *,
         min_width: float,
-        include_ends: Literal["both", "neither", "lower", "upper"],
+        include_endpoints: Literal["both", "neither", "lower", "upper"],
     ):
         scale = interval[1] - interval[0]
         n_widths = arr.shape[-1]
@@ -199,10 +206,16 @@ def __init__(
                 "min_width*n_widths is greater than the interval width, so cannot be "
                 "satisfied."
             )
+
+        if include_endpoints not in ["both", "neither", "lower", "upper"]:
+            raise ValueError(
+                "include_endpoints must be one of 'both', 'neither', 'lower', or 'upper'",
+            )
+
         self.arr = arr
         self.interval = interval
         self.min_width = min_width
-        self.include_ends = include_ends
+        self.include_endpoints = include_endpoints
 
     def unwrap(self) -> Array:
         scale = self.interval[1] - self.interval[0]
@@ -216,9 +229,9 @@ def _unwrap(arr):
             lower, upper = jnp.array([self.interval[0]]), jnp.array([self.interval[1]])
             return jnp.concatenate(
                 [
-                    *([lower] if self.include_ends in ("lower", "both") else []),
+                    *([lower] if self.include_endpoints in ("lower", "both") else []),
                     jnp.cumsum(widths, axis=-1)[:-1] + lower,
-                    *([upper] if self.include_ends in ("upper", "both") else []),
+                    *([upper] if self.include_endpoints in ("upper", "both") else []),
                 ]
             )
 

pyproject.toml

@@ -17,7 +17,7 @@ license = { file = "LICENSE" }
 name = "paramax"
 readme = "README.md"
 requires-python = ">=3.10"
-version = "0.0.3"
+version = "0.0.4"
 
 [project.urls]
 repository = "https://github.com/danielward27/paramax"

tests/test_wrappers.py

@@ -54,7 +54,7 @@ def test_RealToIncreasingOnInterval(case):
         case["arr"],
         (-1, 2),
         min_width=0.1,
-        include_ends=case["include_ends"],
+        include_endpoints=case["include_ends"],
     )
     result = unwrap(real_to_inc)
     assert pytest.approx(case["expected"]) == result
@@ -102,7 +102,7 @@ def test_WeightNormalization():
     "Parameterize-exp": lambda key: Parameterize(jnp.exp, jr.normal(key, (10,))),
     "WeightNormalization": lambda key: WeightNormalization(jr.normal(key, (10, 2))),
     "RealToIncreasingOnInterval": lambda key: RealToIncreasingOnInterval(
-        jnp.zeros(10), (-7, 5), min_width=0.2, include_ends="upper"
+        jnp.zeros(10), (-7, 5), min_width=0.2, include_endpoints="upper"
     ),
 }