Jax kmeans implementation, sklearn removal from main deps

dynamax/hidden_markov_model/models/arhmm.py

@@ -15,6 +15,7 @@
 from dynamax.parameters import ParameterProperties
 from dynamax.types import Scalar
 from dynamax.utils.bijectors import RealToPSDBijector
+from dynamax.utils.cluster import kmeans
 from tensorflow_probability.substrates import jax as tfp
 
 tfd = tfp.distributions
@@ -64,12 +65,8 @@ def initialize(self,
 
         if method.lower() == "kmeans":
             assert emissions is not None, "Need emissions to initialize the model with K-Means!"
-            from sklearn.cluster import KMeans
-            key, subkey = jr.split(key)  # Create a random seed for SKLearn.
-            sklearn_key = jr.randint(subkey, shape=(), minval=0, maxval=2147483647)  # Max int32 value.
-            km = KMeans(self.num_states, random_state=int(sklearn_key)).fit(emissions.reshape(-1, self.emission_dim))
             _emission_weights = jnp.zeros((self.num_states, self.emission_dim, self.emission_dim * self.num_lags))
-            _emission_biases = jnp.array(km.cluster_centers_)
+            _emission_biases = kmeans(emissions.reshape(-1, self.emission_dim), self.num_states, key).centroids
             _emission_covs = jnp.tile(jnp.eye(self.emission_dim)[None, :, :], (self.num_states, 1, 1))
 
         elif method.lower() == "prior":

dynamax/hidden_markov_model/models/gamma_hmm.py

@@ -9,6 +9,7 @@
 from dynamax.hidden_markov_model.models.initial import StandardHMMInitialState, ParamsStandardHMMInitialState
 from dynamax.hidden_markov_model.models.transitions import StandardHMMTransitions, ParamsStandardHMMTransitions
 from dynamax.types import Scalar
+from dynamax.utils.cluster import kmeans
 import optax
 from typing import NamedTuple, Optional, Tuple, Union
 
@@ -67,14 +68,8 @@ def initialize(
 
         if method.lower() == "kmeans":
             assert emissions is not None, "Need emissions to initialize the model with K-Means!"
-            from sklearn.cluster import KMeans
-
-            key, subkey = jr.split(key)  # Create a random seed for SKLearn.
-            sklearn_key = jr.randint(subkey, shape=(), minval=0, maxval=2147483647)  # Max int32 value.
-            km = KMeans(self.num_states, random_state=int(sklearn_key)).fit(emissions.reshape(-1, 1))
-
             _emission_concentrations = jnp.ones((self.num_states,))
-            _emission_rates = jnp.ravel(1.0 / km.cluster_centers_)
+            _emission_rates = jnp.ravel(1.0 / kmeans(emissions.reshape(-1, 1), self.num_states, key).centroids)
 
         elif method.lower() == "prior":
             _emission_concentrations = jnp.ones((self.num_states,))

dynamax/hidden_markov_model/models/gaussian_hmm.py

@@ -21,6 +21,7 @@
 from dynamax.utils.distributions import niw_posterior_update
 from dynamax.utils.bijectors import RealToPSDBijector
 from dynamax.utils.utils import pytree_sum
+from dynamax.utils.cluster import kmeans
 
 
 class ParamsGaussianHMMEmissions(NamedTuple):
@@ -97,12 +98,7 @@ def initialize(self,
         """
         if method.lower() == "kmeans":
             assert emissions is not None, "Need emissions to initialize the model with K-Means!"
-            from sklearn.cluster import KMeans
-            key, subkey = jr.split(key)  # Create a random seed for SKLearn.
-            sklearn_key = jr.randint(subkey, shape=(), minval=0, maxval=2147483647)  # Max int32 value.
-            km = KMeans(self.num_states, random_state=int(sklearn_key)).fit(emissions.reshape(-1, self.emission_dim))
-
-            _emission_means = jnp.array(km.cluster_centers_)
+            _emission_means = kmeans(emissions.reshape(-1, self.emission_dim), self.num_states, key).centroids
             _emission_covs = jnp.tile(jnp.eye(self.emission_dim)[None, :, :], (self.num_states, 1, 1))
 
         elif method.lower() == "prior":
@@ -239,11 +235,7 @@ def initialize(self,
 
         if method.lower() == "kmeans":
             assert emissions is not None, "Need emissions to initialize the model with K-Means!"
-            from sklearn.cluster import KMeans
-            key, subkey = jr.split(key)  # Create a random seed for SKLearn.
-            sklearn_key = jr.randint(subkey, shape=(), minval=0, maxval=2147483647)  # Max int32 value.
-            km = KMeans(self.num_states, random_state=int(sklearn_key)).fit(emissions.reshape(-1, self.emission_dim))
-            _emission_means = jnp.array(km.cluster_centers_)
+            _emission_means = kmeans(emissions.reshape(-1, self.emission_dim), self.num_states, key).centroids
             _emission_scale_diags = jnp.ones((self.num_states, self.emission_dim))
 
         elif method.lower() == "prior":
@@ -383,11 +375,7 @@ def initialize(self,
         """
         if method.lower() == "kmeans":
             assert emissions is not None, "Need emissions to initialize the model with K-Means!"
-            from sklearn.cluster import KMeans
-            key, subkey = jr.split(key)  # Create a random seed for SKLearn.
-            sklearn_key = jr.randint(subkey, shape=(), minval=0, maxval=2147483647)  # Max int32 value.
-            km = KMeans(self.num_states, random_state=int(sklearn_key)).fit(emissions.reshape(-1, self.emission_dim))
-            _emission_means = jnp.array(km.cluster_centers_)
+            _emission_means = kmeans(emissions.reshape(-1, self.emission_dim), self.num_states, key).centroids
             _emission_scales = jnp.ones((self.num_states,))
 
         elif method.lower() == "prior":
@@ -498,11 +486,7 @@ def initialize(self,
         """
         if method.lower() == "kmeans":
             assert emissions is not None, "Need emissions to initialize the model with K-Means!"
-            from sklearn.cluster import KMeans
-            key, subkey = jr.split(key)  # Create a random seed for SKLearn.
-            sklearn_key = jr.randint(subkey, shape=(), minval=0, maxval=2147483647)  # Max int32 value.
-            km = KMeans(self.num_states, random_state=int(sklearn_key)).fit(emissions.reshape(-1, self.emission_dim))
-            _emission_means = jnp.array(km.cluster_centers_)
+            _emission_means = kmeans(emissions.reshape(-1, self.emission_dim), self.num_states, key).centroids
             _emission_cov = jnp.eye(self.emission_dim)
 
         elif method.lower() == "prior":
@@ -663,11 +647,7 @@ def initialize(self,
         """
         if method.lower() == "kmeans":
             assert emissions is not None, "Need emissions to initialize the model with K-Means!"
-            from sklearn.cluster import KMeans
-            key, subkey = jr.split(key)  # Create a random seed for SKLearn.
-            sklearn_key = jr.randint(subkey, shape=(), minval=0, maxval=2147483647)  # Max int32 value.
-            km = KMeans(self.num_states, random_state=int(sklearn_key)).fit(emissions.reshape(-1, self.emission_dim))
-            _emission_means = jnp.array(km.cluster_centers_)
+            _emission_means = kmeans(emissions.reshape(-1, self.emission_dim), self.num_states, key).centroids
             _emission_cov_diag_factors = jnp.ones((self.num_states, self.emission_dim))
             _emission_cov_low_rank_factors = jnp.zeros((self.num_states, self.emission_dim, self.emission_rank))
 

dynamax/hidden_markov_model/models/gmm_hmm.py

@@ -20,6 +20,7 @@
 from dynamax.hidden_markov_model.models.transitions import StandardHMMTransitions, ParamsStandardHMMTransitions
 from dynamax.utils.bijectors import RealToPSDBijector
 from dynamax.utils.utils import pytree_sum
+from dynamax.utils.cluster import kmeans
 from dynamax.types import IntScalar, Scalar
 
 
@@ -114,12 +115,11 @@ def initialize(self,
         """
         if method.lower() == "kmeans":
             assert emissions is not None, "Need emissions to initialize the model with K-Means!"
-            from sklearn.cluster import KMeans
-            key, subkey = jr.split(key)  # Create a random seed for SKLearn.
-            sklearn_key = jr.randint(subkey, shape=(), minval=0, maxval=2147483647)  # Max int32 value.
-            km = KMeans(self.num_states, random_state=int(sklearn_key)).fit(emissions.reshape(-1, self.emission_dim))
             _emission_weights = jnp.ones((self.num_states, self.num_components)) / self.num_components
-            _emission_means = jnp.tile(jnp.array(km.cluster_centers_)[:, None, :], (1, self.num_components, 1))
+            _emission_means = jnp.tile(
+                kmeans(emissions.reshape(-1, self.emission_dim), self.num_states, key).centroids[:, None, :],
+                (1, self.num_components, 1),
+            )
             _emission_covs = jnp.tile(jnp.eye(self.emission_dim), (self.num_states, self.num_components, 1, 1))
 
         elif method.lower() == "prior":
@@ -394,12 +394,11 @@ def initialize(self,
         """
         if method.lower() == "kmeans":
             assert emissions is not None, "Need emissions to initialize the model with K-Means!"
-            from sklearn.cluster import KMeans
-            key, subkey = jr.split(key)  # Create a random seed for SKLearn.
-            sklearn_key = jr.randint(subkey, shape=(), minval=0, maxval=2147483647)  # Max int32 value.
-            km = KMeans(self.num_states, random_state=int(sklearn_key)).fit(emissions.reshape(-1, self.emission_dim))
             _emission_weights = jnp.ones((self.num_states, self.num_components)) / self.num_components
-            _emission_means = jnp.tile(jnp.array(km.cluster_centers_)[:, None, :], (1, self.num_components, 1))
+            _emission_means = jnp.tile(
+                kmeans(emissions.reshape(-1, self.emission_dim), self.num_states, key).centroids[:, None, :],
+                (1, self.num_components, 1),
+            )
             _emission_scale_diags = jnp.ones((self.num_states, self.num_components, self.emission_dim))
 
         elif method.lower() == "prior":

dynamax/hidden_markov_model/models/linreg_hmm.py

@@ -16,6 +16,7 @@
 from dynamax.types import Scalar
 from dynamax.utils.utils import pytree_sum
 from dynamax.utils.bijectors import RealToPSDBijector
+from dynamax.utils.cluster import kmeans
 
 tfd = tfp.distributions
 tfb = tfp.bijectors
@@ -80,12 +81,8 @@ def initialize(self,
         """
         if method.lower() == "kmeans":
             assert emissions is not None, "Need emissions to initialize the model with K-Means!"
-            from sklearn.cluster import KMeans
-            key, subkey = jr.split(key)  # Create a random seed for SKLearn.
-            sklearn_key = jr.randint(subkey, shape=(), minval=0, maxval=2147483647)  # Max int32 value.
-            km = KMeans(self.num_states, random_state=int(sklearn_key)).fit(emissions.reshape(-1, self.emission_dim))
             _emission_weights = jnp.zeros((self.num_states, self.emission_dim, self.input_dim))
-            _emission_biases = jnp.array(km.cluster_centers_)
+            _emission_biases = kmeans(emissions.reshape(-1, self.emission_dim), self.num_states, key).centroids
             _emission_covs = jnp.tile(jnp.eye(self.emission_dim)[None, :, :], (self.num_states, 1, 1))
 
         elif method.lower() == "prior":

dynamax/hidden_markov_model/models/logreg_hmm.py

@@ -14,6 +14,7 @@
 from dynamax.hidden_markov_model.models.initial import StandardHMMInitialState, ParamsStandardHMMInitialState
 from dynamax.hidden_markov_model.models.transitions import StandardHMMTransitions, ParamsStandardHMMTransitions
 from dynamax.types import IntScalar, Scalar
+from dynamax.utils.cluster import kmeans
 
 
 class ParamsLogisticRegressionHMMEmissions(NamedTuple):
@@ -80,16 +81,19 @@ def initialize(self,
         if method.lower() == "kmeans":
             assert emissions is not None, "Need emissions to initialize the model with K-Means!"
             assert inputs is not None, "Need inputs to initialize the model with K-Means!"
-            from sklearn.cluster import KMeans
-
             flat_emissions = emissions.reshape(-1,)
             flat_inputs = inputs.reshape(-1, self.input_dim)
-            key, subkey = jr.split(key)  # Create a random seed for SKLearn.
-            sklearn_key = jr.randint(subkey, shape=(), minval=0, maxval=2147483647)  # Max int32 value.
-            km = KMeans(self.num_states, random_state=int(sklearn_key)).fit(flat_inputs)
+
+            assignments = kmeans(flat_inputs, self.num_states, key).assignments
             _emission_weights = jnp.zeros((self.num_states, self.input_dim))
-            _emission_biases = jnp.array([tfb.Sigmoid().inverse(flat_emissions[km.labels_ == k].mean())
-                                          for k in range(self.num_states)])
+            # Both guards keep the logit below finite: an empty cluster has an undefined
+            # mean, and an all-0s or all-1s cluster would otherwise logit to -/+ infinity.
+            cluster_means = jnp.array(
+                [jnp.mean(flat_emissions, where=(assignments == k)) for k in range(self.num_states)]
+            )
+            cluster_means = jnp.where(jnp.isnan(cluster_means), flat_emissions.mean(), cluster_means)
+            eps = jnp.finfo(cluster_means.dtype).eps
+            _emission_biases = tfb.Sigmoid().inverse(jnp.clip(cluster_means, eps, 1.0 - eps))
 
         elif method.lower() == "prior":
             # TODO: Use an MNIW prior

dynamax/hidden_markov_model/models/test_models.py

@@ -1,6 +1,7 @@
 """Tests for the HMM models."""
 
 import dynamax.hidden_markov_model as models
+import jax
 import jax.numpy as jnp
 import jax.random as jr
 import pytest
@@ -169,3 +170,127 @@ def test_sample_and_fit_arhmm():
 #     test_cov = test_hmm.emission_covariance_matrices.value
 #     assert jnp.alltrue(test_cov.shape == (10, 2, 2))
 #     assert jnp.allclose(jnp.linalg.norm(test_cov-refr_cov, axis=-1), 0., atol=1)
+
+
+# Models whose emissions support kmeans initialization. BernoulliHMM, CategoricalHMM,
+# MultinomialHMM and PoissonHMM deliberately raise NotImplementedError instead.
+KMEANS_CONFIGS = [
+    # (cls, kwargs, inputs, init_takes_inputs)
+    (models.GammaHMM, dict(num_states=4), None, False),
+    (models.GaussianHMM, dict(num_states=4, emission_dim=3), None, False),
+    (models.DiagonalGaussianHMM, dict(num_states=4, emission_dim=3), None, False),
+    (models.SphericalGaussianHMM, dict(num_states=4, emission_dim=3), None, False),
+    (models.SharedCovarianceGaussianHMM, dict(num_states=4, emission_dim=3), None, False),
+    (models.LowRankGaussianHMM, dict(num_states=4, emission_dim=3, emission_rank=1), None, False),
+    (models.GaussianMixtureHMM, dict(num_states=4, num_components=2, emission_dim=3), None, False),
+    (models.DiagonalGaussianMixtureHMM, dict(num_states=4, num_components=2, emission_dim=3), None, False),
+    (models.LinearRegressionHMM, dict(num_states=3, emission_dim=3, input_dim=5),
+     jr.normal(jr.PRNGKey(0), (NUM_TIMESTEPS, 5)), False),
+    (models.LogisticRegressionHMM, dict(num_states=4, input_dim=5),
+     jr.normal(jr.PRNGKey(0), (NUM_TIMESTEPS, 5)), True),
+]
+
+
+@pytest.mark.parametrize(["cls", "kwargs", "inputs", "init_takes_inputs"], KMEANS_CONFIGS)
+def test_initialize_kmeans(cls, kwargs, inputs, init_takes_inputs):
+    """Test that kmeans initialization produces finite, fittable parameters."""
+    hmm = cls(**kwargs)
+    key1, key2, key3 = jr.split(jr.PRNGKey(42), 3)
+    params, _ = hmm.initialize(key1)
+    _, emissions = hmm.sample(params, key2, num_timesteps=NUM_TIMESTEPS, inputs=inputs)
+
+    # Only LogisticRegressionHMM clusters its inputs, so only it accepts them here.
+    init_kwargs = dict(emissions=emissions)
+    if init_takes_inputs:
+        init_kwargs["inputs"] = inputs
+    km_params, km_props = hmm.initialize(key3, method="kmeans", **init_kwargs)
+
+    # No parameter may come back NaN or infinite; that was the failure mode of an
+    # unguarded empty cluster.
+    for leaf in jax.tree_util.tree_leaves(km_params):
+        assert jnp.all(jnp.isfinite(leaf))
+
+    # The initialization must be usable: EM from it improves monotonically.
+    _, lps = hmm.fit_em(km_params, km_props, emissions, inputs=inputs, num_iters=3, verbose=False)
+    assert monotonically_increasing(lps, atol=1e-2, rtol=1e-2)
+
+
+def test_initialize_kmeans_arhmm():
+    """Test that kmeans initialization works for a LinearAutoregressiveHMM."""
+    arhmm = models.LinearAutoregressiveHMM(num_states=4, emission_dim=2, num_lags=1)
+    key1, key2, key3 = jr.split(jr.PRNGKey(42), 3)
+    params, _ = arhmm.initialize(key1)
+    _, emissions = arhmm.sample(params, key2, num_timesteps=NUM_TIMESTEPS)
+
+    km_params, _ = arhmm.initialize(key3, method="kmeans", emissions=emissions)
+
+    for leaf in jax.tree_util.tree_leaves(km_params):
+        assert jnp.all(jnp.isfinite(leaf))
+    assert km_params.emissions.biases.shape == (4, 2)
+
+
+def test_logreg_hmm_kmeans_finite_bias_with_saturated_cluster():
+    """A cluster whose assigned emissions are all 0 (or all 1) is not the same as an
+    empty cluster, but it drives the cluster mean to exactly 0.0 or 1.0. Feeding that
+    straight into the logit used to produce +/-inf biases even though the existing
+    NaN guard (for genuinely empty clusters) passed. Construct inputs as two
+    well-separated blobs so kmeans reliably keeps them apart, and make one blob's
+    emissions uniformly 0 and the other uniformly 1, so the cluster means saturate
+    at both 0.0 and 1.0 and exercise both the lower and upper clip bounds.
+    """
+    from dynamax.hidden_markov_model.models.logreg_hmm import LogisticRegressionHMMEmissions
+
+    emission_component = LogisticRegressionHMMEmissions(num_states=2, input_dim=2)
+    inputs = jnp.concatenate([
+        jnp.tile(jnp.array([1000.0, 0.0]), (10, 1)),
+        jnp.tile(jnp.array([-1000.0, 0.0]), (10, 1)),
+    ], axis=0)
+    emissions = jnp.concatenate([
+        jnp.zeros(10),
+        jnp.ones(10),
+    ], axis=0)
+
+    for seed in range(10):
+        params, _ = emission_component.initialize(
+            jr.PRNGKey(seed), method="kmeans", emissions=emissions, inputs=inputs)
+        assert jnp.all(jnp.isfinite(params.biases)), f"seed {seed} produced a non-finite bias"
+        # Sanity check that this seed really did hit both the low- and high-saturated
+        # cluster cases (the clip's eps floor and its 1 - eps ceiling).
+        assert jnp.any(params.biases < 0)
+        assert jnp.any(params.biases > 0)
+
+
+def test_initialize_kmeans_is_jax_transformable():
+    """Kmeans initialization must stay compatible with `jax.jit` and `jax.vmap`.
+
+    `LogisticRegressionHMMEmissions.initialize`'s kmeans branch computes each
+    cluster's emission mean with `jnp.mean(flat_emissions, where=(assignments == k))`.
+    The seemingly-equivalent `flat_emissions[assignments == k].mean()` also passes
+    eagerly, but it boolean-mask-indexes with a *traced* array, which produces a
+    variable-shaped intermediate. That is illegal under `jax.jit`/`jax.vmap` (it
+    raises `NonConcreteBooleanIndexError`), even though nothing catches it outside
+    a JAX transformation. If this test starts failing with that error, the fix is
+    to restore the `where=` form in the kmeans branch of `initialize` -- not to
+    delete this test.
+    """
+    num_states, input_dim = 4, 5
+    hmm = models.LogisticRegressionHMM(num_states=num_states, input_dim=input_dim)
+    key1, key2 = jr.split(jr.PRNGKey(0), 2)
+    params, _ = hmm.initialize(key1)
+    inputs = jr.normal(key2, (NUM_TIMESTEPS, input_dim))
+    _, emissions = hmm.sample(params, jr.PRNGKey(1), num_timesteps=NUM_TIMESTEPS, inputs=inputs)
+
+    def init_kmeans(key, emissions, inputs):
+        km_params, _ = hmm.initialize(key, method="kmeans", emissions=emissions, inputs=inputs)
+        return km_params
+
+    jitted_params = jax.jit(init_kmeans)(jr.PRNGKey(2), emissions, inputs)
+    assert jnp.all(jnp.isfinite(jitted_params.emissions.biases))
+
+    batch_size = 3
+    batch_keys = jr.split(jr.PRNGKey(3), batch_size)
+    batch_emissions = jnp.stack([emissions] * batch_size)
+    batch_inputs = jnp.stack([inputs] * batch_size)
+    vmapped_params = vmap(init_kmeans)(batch_keys, batch_emissions, batch_inputs)
+    assert vmapped_params.emissions.biases.shape == (batch_size, num_states)
+    assert jnp.all(jnp.isfinite(vmapped_params.emissions.biases))