minor bug fix on initial sample noon_finite_guess

kazewong · Sep 16, 2024 · aef6124 · aef6124
1 parent 24235c2
commit aef6124
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/example/GW150914_IMRPhenomD.py b/example/GW150914_IMRPhenomD.py
@@ -117,7 +117,7 @@
     n_epochs=n_epochs,
     learning_rate=learning_rate,
     n_max_examples=30000,
-    n_flow_samples=100000,
+    n_flow_sample=100000,
     momentum=0.9,
     batch_size=30000,
     use_global=True,

diff --git a/src/jimgw/jim.py b/src/jimgw/jim.py
@@ -107,7 +107,7 @@ def sample(self, key: PRNGKeyArray, initial_position: Array = jnp.array([])):
             initial_position = jnp.zeros((self.sampler.n_chains, self.prior.n_dim)) + jnp.nan
 
             while not jax.tree.reduce(jnp.logical_and, jax.tree.map(lambda x: jnp.isfinite(x), initial_position)).all():
-                non_finite_index = jnp.any(~jax.tree.reduce(jnp.logical_and, jax.tree.map(lambda x: jnp.isfinite(x), initial_position)),axis=1)
+                non_finite_index = jnp.where(jnp.any(~jax.tree.reduce(jnp.logical_and, jax.tree.map(lambda x: jnp.isfinite(x), initial_position)),axis=1))[0]
 
                 key, subkey = jax.random.split(key)
                 guess = self.prior.sample(subkey, self.sampler.n_chains)

diff --git a/test/integration/test_GW150914_D.py b/test/integration/test_GW150914_D.py
@@ -119,7 +119,7 @@
     n_epochs=n_epochs,
     learning_rate=learning_rate,
     n_max_examples=30,
-    n_flow_samples=100,
+    n_flow_sample=100,
     momentum=0.9,
     batch_size=100,
     use_global=True,