Merge branch 'feat/beartype-claw-cleanup' into feat/beartype-claw-extend

Author: hmgaudecker

src/lcm/_beartype_conf.py

@@ -6,9 +6,7 @@
 
 """
 
-from collections.abc import Callable
-
-from beartype import BeartypeConf, BeartypeStrategy, beartype
+from beartype import BeartypeConf, BeartypeStrategy
 
 from lcm.exceptions import (
     CategoricalDefinitionError,
@@ -34,24 +32,6 @@ def _conf(exc: type[Exception]) -> BeartypeConf:
     )
 
 
-def beartype_init[C](conf: BeartypeConf) -> Callable[[type[C]], type[C]]:
-    """Class decorator that beartype-checks `__init__` only.
-
-    Bare `@beartype` on a class wraps every method, which surfaces
-    annotation drift in helpers like `compute_gridpoints(**kwargs: float)`
-    where runtime kwargs are actually JAX arrays. Restricting decoration
-    to `__init__` keeps the perimeter check (parameter types at
-    construction) without policing every method's runtime types.
-
-    """
-
-    def deco(cls: type[C]) -> type[C]:
-        cls.__init__ = beartype(conf=conf)(cls.__init__)  # ty: ignore[invalid-assignment]
-        return cls
-
-    return deco
-
-
 # Used on `Regime` and `MarkovTransition`.
 REGIME_CONF = _conf(RegimeInitializationError)
 

src/lcm/grids/continuous.py

@@ -5,7 +5,6 @@
 
 import jax.numpy as jnp
 
-from lcm._beartype_conf import GRID_CONF, beartype_init
 from lcm.dtypes import canonical_float_dtype
 from lcm.exceptions import GridInitializationError, format_messages
 from lcm.grids import coordinates as grid_coordinates
@@ -100,7 +99,6 @@ def replace(self, **kwargs: float) -> UniformContinuousGrid:
             ) from e
 
 
-@beartype_init(GRID_CONF)
 class LinSpacedGrid(UniformContinuousGrid):
     """A linearly spaced grid of continuous values.
 
@@ -126,7 +124,6 @@ def get_coordinate(self, value: FloatND) -> FloatND:
         )
 
 
-@beartype_init(GRID_CONF)
 class LogSpacedGrid(UniformContinuousGrid):
     """A logarithmically spaced grid of continuous values.
 
@@ -209,7 +206,6 @@ def _init_uniform_grid(
     object.__setattr__(grid, "distributed", distributed)
 
 
-@beartype_init(GRID_CONF)
 @dataclass(frozen=True, kw_only=True, init=False)
 class IrregSpacedGrid(ContinuousGrid):
     """A grid of continuous values at irregular (user-specified) points.

src/lcm/grids/discrete.py

@@ -1,13 +1,11 @@
 import jax.numpy as jnp
 
-from lcm._beartype_conf import GRID_CONF, beartype_init
 from lcm.grids.base import Grid
 from lcm.grids.categorical import _validate_discrete_grid
 from lcm.typing import Int1D
 from lcm.utils.containers import get_field_names_and_values
 
 
-@beartype_init(GRID_CONF)
 class DiscreteGrid(Grid):
     """A discrete grid defining the outcome space of a categorical variable.
 

src/lcm/grids/piecewise.py

@@ -4,7 +4,6 @@
 import jax.numpy as jnp
 import portion
 
-from lcm._beartype_conf import GRID_CONF, beartype_init
 from lcm.exceptions import GridInitializationError, format_messages
 from lcm.grids import coordinates as grid_coordinates
 from lcm.grids.continuous import ContinuousGrid
@@ -44,7 +43,6 @@ def __init__(
         object.__setattr__(self, "n_points", jnp.int32(n_points))
 
 
-@beartype_init(GRID_CONF)
 @dataclass(frozen=True, kw_only=True)
 class PiecewiseLinSpacedGrid(ContinuousGrid):
     """A piecewise linearly spaced grid with multiple segments.
@@ -108,7 +106,6 @@ def get_coordinate(self, value: FloatND) -> FloatND:
         return self._cumulative_offsets[piece_idx] + local_coord
 
 
-@beartype_init(GRID_CONF)
 @dataclass(frozen=True, kw_only=True)
 class PiecewiseLogSpacedGrid(ContinuousGrid):
     """A piecewise logarithmically spaced grid with multiple segments.

src/lcm/shocks/ar1.py

@@ -7,7 +7,6 @@
 import jax.numpy as jnp
 from jax.scipy.stats.norm import cdf
 
-from lcm._beartype_conf import GRID_CONF, beartype_init
 from lcm.shocks._base import (
     _gauss_hermite_normal,
     _mixture_cdf,
@@ -30,7 +29,6 @@ def draw_shock(
     ) -> ScalarFloat: ...
 
 
-@beartype_init(GRID_CONF)
 @dataclass(frozen=True, kw_only=True)
 class Tauchen(_ShockGridAR1):
     r"""AR(1) shock discretized via Tauchen (1986).
@@ -130,7 +128,6 @@ def draw_shock(
         )
 
 
-@beartype_init(GRID_CONF)
 @dataclass(frozen=True, kw_only=True)
 class Rouwenhorst(_ShockGridAR1):
     r"""AR(1) shock discretized via Rouwenhorst (1995).
@@ -199,7 +196,6 @@ def draw_shock(
         )
 
 
-@beartype_init(GRID_CONF)
 @dataclass(frozen=True, kw_only=True)
 class TauchenNormalMixture(_ShockGridAR1):
     r"""AR(1) shock with mixture-of-normals innovations, discretized via Tauchen.

src/lcm/shocks/iid.py

@@ -6,7 +6,6 @@
 import jax.numpy as jnp
 from jax.scipy.stats.norm import cdf
 
-from lcm._beartype_conf import GRID_CONF, beartype_init
 from lcm.shocks._base import (
     _gauss_hermite_normal,
     _mixture_cdf,
@@ -28,7 +27,6 @@ def draw_shock(
     ) -> ScalarFloat: ...
 
 
-@beartype_init(GRID_CONF)
 @dataclass(frozen=True, kw_only=True)
 class Uniform(_ShockGridIID):
     r"""Discretized iid uniform shock: $U(\text{start}, \text{stop})$.
@@ -63,7 +61,6 @@ def draw_shock(
         )
 
 
-@beartype_init(GRID_CONF)
 @dataclass(frozen=True, kw_only=True)
 class Normal(_ShockGridIID):
     r"""Discretized iid normal shock: $N(\mu_\varepsilon, \sigma_\varepsilon^2)$.
@@ -138,7 +135,6 @@ def draw_shock(
         return params["mu"] + params["sigma"] * jax.random.normal(key=key)
 
 
-@beartype_init(GRID_CONF)
 @dataclass(frozen=True, kw_only=True)
 class LogNormal(_ShockGridIID):
     r"""Discretized iid log-normal shock: $\ln X \sim N(\mu, \sigma^2)$."""
@@ -204,7 +200,6 @@ def draw_shock(
         return jnp.exp(params["mu"] + params["sigma"] * jax.random.normal(key=key))
 
 
-@beartype_init(GRID_CONF)
 @dataclass(frozen=True, kw_only=True)
 class NormalMixture(_ShockGridIID):
     r"""Discretized IID normal-mixture shock.