fixes PRNGKey deprecation (#15)

* fixes #14

PRNGKey Array dtype

* switch PRNGKey to key

* drop 3.8 support

* fix lint

---------

Co-authored-by: Omar Younis <omar.younis98@gmail.com>

Author: jeffhsu3

.github/workflows/build.yml

@@ -15,7 +15,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ['3.8', '3.9', '3.10', '3.11']
+        python-version: ['3.9', '3.10', '3.11', '3.12']
     steps:
     - uses: actions/checkout@v3
     - name: Set up Python ${{ matrix.python-version }}

chromax/functional.py

@@ -13,7 +13,7 @@
 def cross(
     parents: Parents["n"],
     recombination_vec: Float[Array, N_MARKERS],
-    random_key: jax.random.PRNGKeyArray,
+    random_key: jax.Array,
     mutation_probability: float = 0.0,
 ) -> Population["n"]:
     """Main function that computes crosses from a list of parents.
@@ -25,8 +25,8 @@ def cross(
     :param recombination_vec: array of m probabilities.
         The i-th value represent the probability to recombine before the marker i.
     :type recombination_vec:
-    :param random_key: JAX PRNGKey, for reproducibility purpose.
-    :type random_key: jax.random.PRNGKeyArray
+    :param random_key: JAX random key, for reproducibility purpose.
+    :type random_key: jax.Array
     :param mutation_probability: The probability of having a mutation in a marker.
     :type mutation_probability: float
     :return: offspring population of shape (n, m, d).
@@ -43,7 +43,7 @@ def cross(
         >>> rec_vec = np.full((n_chr, chr_len), 1.5 / chr_len)
         >>> rec_vec[:, 0] = 0.5  # equal probability on starting haploid
         >>> rec_vec = rec_vec.flatten()
-        >>> random_key = jax.random.PRNGKey(42)
+        >>> random_key = jax.random.key(42)
         >>> f2 = functional.cross(parents, rec_vec, random_key)
         >>> f2.shape
         (50, 1000, 2)
@@ -52,7 +52,7 @@ def cross(
     random_keys = jax.random.split(
         random_key, num=2 * len(parents) * 2 * parents.shape[3]
     )
-    random_keys = random_keys.reshape(2, len(parents), 2, parents.shape[3], 2)
+    random_keys = random_keys.reshape(2, len(parents), 2, parents.shape[3])
     cross_random_key, mutate_random_key = random_keys
 
     offsprings = _cross(
@@ -71,8 +71,8 @@ def cross(
 def _cross(
     parent: Individual,
     recombination_vec: Float[Array, N_MARKERS],
-    cross_random_key: jax.random.PRNGKeyArray,
-    mutate_random_key: jax.random.PRNGKeyArray,
+    cross_random_key: jax.Array,
+    mutate_random_key: jax.Array,
     mutation_probability: float,
 ) -> Haploid:
     return _meiosis(
@@ -88,7 +88,7 @@ def double_haploid(
     population: Population["n"],
     n_offspring: int,
     recombination_vec: Float[Array, N_MARKERS],
-    random_key: jax.random.PRNGKeyArray,
+    random_key: jax.Array,
     mutation_probability: float = 0.0,
 ) -> Population["n n_offspring"]:
     """Computes the double haploid of the input population.
@@ -100,8 +100,8 @@ def double_haploid(
     :param recombination_vec: array of m probabilities.
         The i-th value represent the probability to recombine before the marker i.
     :type recombination_vec: ndarray
-    :param random_key: array of n PRNGKey, one for each individual.
-    :type random_key: jax.random.PRNGKeyArray
+    :param random_key: JAX random key, for reproducibility purpose.
+    :type random_key: jax.Array
     :param mutation_probability: The probability of having a mutation in a marker.
     :type mutation_probability: float
     :return: output population of shape (n, n_offspring, m, d).
@@ -118,15 +118,15 @@ def double_haploid(
         >>> rec_vec = np.full((n_chr, chr_len), 1.5 / chr_len)
         >>> rec_vec[:, 0] = 0.5  # equal probability on starting haploid
         >>> rec_vec = rec_vec.flatten()
-        >>> random_key = jax.random.PRNGKey(42)
+        >>> random_key = jax.random.key(42)
         >>> dh = functional.double_haploid(f1, 10, rec_vec, random_key)
         >>> dh.shape
         (50, 10, 1000, 2)
     """
     population = population.reshape(*population.shape[:2], -1, 2)
     keys = jax.random.split(
         random_key, num=2 * len(population) * n_offspring * population.shape[2]
-    ).reshape(2, len(population), n_offspring, population.shape[2], 2)
+    ).reshape(2, len(population), n_offspring, population.shape[2])
     cross_random_key, mutate_random_key = keys
     haploids = _double_haploid(
         population,
@@ -145,8 +145,8 @@ def double_haploid(
 def _double_haploid(
     individual: Individual,
     recombination_vec: Float[Array, N_MARKERS],
-    cross_random_key: jax.random.PRNGKeyArray,
-    mutate_random_key: jax.random.PRNGKeyArray,
+    cross_random_key: jax.Array,
+    mutate_random_key: jax.Array,
     mutation_probability: float,
 ) -> Haploid:
     return _meiosis(
@@ -165,8 +165,8 @@ def _double_haploid(
 def _meiosis(
     individual: Individual,
     recombination_vec: Float[Array, N_MARKERS],
-    cross_random_key: jax.random.PRNGKeyArray,
-    mutate_random_key: jax.random.PRNGKeyArray,
+    cross_random_key: jax.Array,
+    mutate_random_key: jax.Array,
     mutation_probability: float,
 ) -> Haploid:
     samples = jax.random.uniform(cross_random_key, shape=recombination_vec.shape)

chromax/simulator.py

@@ -171,7 +171,7 @@ def set_seed(self, seed: int):
         :param seed: random seed.
         :type seed: int
         """
-        self.random_key = jax.random.PRNGKey(seed)
+        self.random_key = jax.random.key(seed)
 
     def load_population(self, file_name: Union[Path, str]) -> Population["n"]:
         """Load a population from file.
@@ -271,7 +271,7 @@ def differentiable_cross_func(self) -> Callable:
             >>> f1 = simulator.load_population(sample_data.genome)
             >>> weights = np.random.uniform(size=(10, len(f1), 2))
             >>> weights /= weights.sum(axis=1, keepdims=True)
-            >>> random_key = jax.random.PRNGKey(42)
+            >>> random_key = jax.random.key(42)
             >>> grad_value = grad_f(f1, weights, random_key)
             >>> grad_value.shape
             (10, 371, 2)
@@ -284,7 +284,7 @@ def differentiable_cross_func(self) -> Callable:
         def diff_cross_f(
             population: Population["n"],
             cross_weights: Float[Array, "m n 2"],
-            random_key: jax.random.PRNGKeyArray,
+            random_key: jax.Array,
         ) -> Population["m"]:
             population = population.reshape(*population.shape[:-1], -1, 2)
             keys_shape = len(cross_weights), len(population), 2, population.shape[-2]

pyproject.toml

@@ -8,20 +8,19 @@ build-backend = "setuptools.build_meta"
 name = "chromax"
 description = "Breeding simulator based on JAX"
 readme = "README.md"
-requires-python = ">= 3.8"
+requires-python = ">= 3.9"
 authors = [{ name = "Omar G. Younis", email = "omar.younis98@gmail.com" }]
 license = { text = "BSD-3-Clause" }
 keywords = ["Breeding", "simulator", "JAX", "chromosome", "genetics", "bioinformatics",]
 classifiers = [
     "Development Status :: 4 - Beta",
     "Programming Language :: Python :: 3",
-    "Programming Language :: Python :: 3.8",
     "Programming Language :: Python :: 3.9",
     "Programming Language :: Python :: 3.10",
     "Programming Language :: Python :: 3.11",
     'Intended Audience :: Science/Research',
 ]
-dependencies = ["numpy", "pandas", "jax", "jaxlib", "jaxtyping"]
+dependencies = ["numpy", "pandas", "jax>=0.4.16", "jaxlib>=0.4.16", "jaxtyping"]
 dynamic = ["version"]
 
 [project.optional-dependencies]

tests/test_functional.py

@@ -15,7 +15,7 @@ def test_cross(idx):
     parents = np.random.choice([False, True], size=parents_shape)
     rec_vec = np.zeros(n_markers)
     rec_vec[0] = idx
-    random_key = jax.random.PRNGKey(42)
+    random_key = jax.random.key(42)
     new_pop = functional.cross(parents, rec_vec, random_key)
 
     for i in range(ploidy):
@@ -28,7 +28,7 @@ def test_double_haploid():
     pop_shape = (50, n_chr * chr_len, ploidy)
     f1 = np.random.choice([False, True], size=pop_shape)
     rec_vec = np.full((n_chr * chr_len,), 1.5 / chr_len)
-    random_key = jax.random.PRNGKey(42)
+    random_key = jax.random.key(42)
     dh = functional.double_haploid(f1, n_offspring, rec_vec, random_key)
     assert dh.shape == (len(f1), n_offspring, n_chr * chr_len, ploidy)
 
@@ -62,7 +62,7 @@ def test_cross_mutation():
     rec_vec = np.full((n_markers,), 1.5e-2)
     cross = functional.cross
 
-    random_key = jax.random.PRNGKey(42)
+    random_key = jax.random.key(42)
     assert np.all(cross(zeros_pop, rec_vec, random_key) == 0)
     assert np.all(cross(zeros_pop, rec_vec, random_key, 1) == 1)
     mutated_pop = cross(zeros_pop, rec_vec, random_key, 0.5)
@@ -83,7 +83,7 @@ def test_dh_mutation():
     rec_vec = np.full((n_markers,), 1.5e-2)
     dh = functional.double_haploid
 
-    random_key = jax.random.PRNGKey(42)
+    random_key = jax.random.key(42)
     assert np.all(dh(zeros_pop, 10, rec_vec, random_key) == 0)
     assert np.all(dh(zeros_pop, 10, rec_vec, random_key, 1) == 1)
     mutated_pop = dh(zeros_pop, 10, rec_vec, random_key, 0.5)