add parent-ids info on random_crosses

README.md

@@ -18,25 +18,26 @@ from chromax.sample_data import genetic_map, genome
 
 simulator = Simulator(genetic_map=genetic_map)
 f1 = simulator.load_population(genome)
-f2 = simulator.random_crosses(f1, n_crosses=10, n_offspring=20)
+f2, parent_ids = simulator.random_crosses(f1, n_crosses=10, n_offspring=20)
 ```
 
 ## Citing ChromaX
 
 ```bibtex
 
-@article{Younis2023.05.29.542709,
-	abstract = {ChromaX is a Python library that enables the simulation of genetic recombination, genomic estimated breeding value calculations, and selection processes. By utilizing GPU processing, it can perform these simulations up to two orders of magnitude faster than existing tools with standard hardware. This offers breeders and scientists new opportunities to simulate genetic gain and optimize breeding schemes.},
-	author = {Omar G. Younis and Matteo Turchetta and Daniel Ariza Suarez and Steven Yates and Bruno Studer and Ioannis N. Athanasiadis and Andreas Krause and Joachim M. Buhmann and Luca Corinzia},
-	doi = {10.1101/2023.05.29.542709},
-	elocation-id = {2023.05.29.542709},
-	eprint = {https://www.biorxiv.org/content/early/2023/05/31/2023.05.29.542709.1.full.pdf},
-	journal = {bioRxiv},
-	publisher = {Cold Spring Harbor Laboratory},
-	title = {ChromaX: a fast and scalable breeding program simulator},
-	url = {https://www.biorxiv.org/content/early/2023/05/31/2023.05.29.542709.1},
-	year = {2023},
-	bdsk-url-1 = {https://www.biorxiv.org/content/early/2023/05/31/2023.05.29.542709.1},
-	bdsk-url-2 = {https://doi.org/10.1101/2023.05.29.542709}
+@article{10.1093/bioinformatics/btad691,
+    author = {Younis, Omar G and Turchetta, Matteo and Ariza Suarez, Daniel and Yates, Steven and Studer, Bruno and Athanasiadis, Ioannis N and Krause, Andreas and Buhmann, Joachim M and Corinzia, Luca},
+    title = "{ChromaX: a fast and scalable breeding program simulator}",
+    journal = {Bioinformatics},
+    volume = {39},
+    number = {12},
+    pages = {btad691},
+    year = {2023},
+    month = {11},
+    abstract = "{ChromaX is a Python library that enables the simulation of genetic recombination, genomic estimated breeding value calculations, and selection processes. By utilizing GPU processing, it can perform these simulations up to two orders of magnitude faster than existing tools with standard hardware. This offers breeders and scientists new opportunities to simulate genetic gain and optimize breeding schemes.The documentation is available at https://chromax.readthedocs.io. The code is available at https://github.com/kora-labs/chromax.}",
+    issn = {1367-4811},
+    doi = {10.1093/bioinformatics/btad691},
+    url = {https://doi.org/10.1093/bioinformatics/btad691},
+    eprint = {https://academic.oup.com/bioinformatics/article-pdf/39/12/btad691/54143193/btad691.pdf},
 }
 ```

chromax/simulator.py

@@ -54,7 +54,7 @@ class Simulator:
         >>> from chromax import Simulator, sample_data
         >>> simulator = Simulator(genetic_map=sample_data.genetic_map)
         >>> f1 = simulator.load_population(sample_data.genome)
-        >>> f2 = simulator.random_crosses(f1, n_crosses=10, n_offspring=20)
+        >>> f2, _ = simulator.random_crosses(f1, n_crosses=10, n_offspring=20)
         >>> f2.shape
         (10, 20, 9839, 2)
     """
@@ -198,7 +198,7 @@ def save_population(self, population: Population["n"], file_name: Union[Path, st
             >>> from chromax import Simulator, sample_data
             >>> simulator = Simulator(genetic_map=sample_data.genetic_map)
             >>> f1 = simulator.load_population(sample_data.genome)
-            >>> f2 = simulator.random_crosses(f1, n_crosses=10, n_offspring=20)
+            >>> f2, _ = simulator.random_crosses(f1, n_crosses=10, n_offspring=20)
             >>> simulator.save_population(f2, "pop_file")
         """
         np.save(file_name, population, allow_pickle=False)
@@ -372,16 +372,19 @@ def random_crosses(
             The default value is 1.
         :type n_offspring: int
 
-        :return: output population of shape (n_crosses, n_offspring, m, d).
-        :rtype: ndarray
+        :return: output population of shape (n_crosses, n_offspring, m, d)
+            and parent indices of shape (n_crosses, 2) of performed crosses.
+        :rtype: tuple of two ndarrays
 
         :Example:
             >>> from chromax import Simulator, sample_data
             >>> simulator = Simulator(genetic_map=sample_data.genetic_map)
             >>> f1 = simulator.load_population(sample_data.genome)
-            >>> f2 = simulator.random_crosses(f1, 100, n_offspring=10)
+            >>> f2, parent_ids = simulator.random_crosses(f1, 100, n_offspring=10)
             >>> f2.shape
             (100, 10, 9839, 2)
+            >>> parent_ids.shape
+            (100, 2)
         """
         all_indices = np.arange(len(population))
         diallel_indices = self._diallel_indices(all_indices)
@@ -392,14 +395,14 @@ def random_crosses(
         random_select_idx = jax.random.choice(
             split_key, len(diallel_indices), shape=(n_crosses,), replace=False
         )
-        cross_indices = diallel_indices[random_select_idx]
+        parent_indices = diallel_indices[random_select_idx]
 
-        cross_indices = np.repeat(cross_indices, n_offspring, axis=0)
+        cross_indices = np.repeat(parent_indices, n_offspring, axis=0)
         out = self.cross(population[cross_indices])
         out = out.reshape(n_crosses, n_offspring, *out.shape[1:])
         if n_offspring == 1:
             out = out.squeeze(1)
-        return out
+        return out, parent_indices
 
     def select(
         self,

docs/tutorials/wheat_bp.rst

@@ -29,7 +29,7 @@ We start the breeding program by performing some random crosses that produce a n
 
 .. code-block:: python
 
-    f1 = simulator.random_crosses(f0, 100)
+    f1, _ = simulator.random_crosses(f0, 100)
     dh_lines = simulator.double_haploid(f1, n_offspring=100)
 
 In this way we obtain 100 lines, each containing 100 plants.

examples/time_wheat_bp.py

@@ -12,7 +12,7 @@ def wheat_schema(
     germplasm,
     factor=1,
 ):
-    f1 = simulator.random_crosses(germplasm, 200 * factor)
+    f1, _ = simulator.random_crosses(germplasm, 200 * factor)
 
     dh_lines = simulator.double_haploid(f1, n_offspring=100)
     # Use the following instead on M1 with 20x factor to avoid trashing

examples/wheat_bp.py

@@ -14,7 +14,7 @@
 def wheat_schema(
     simulator: Simulator, germplasm: Population["50"]
 ) -> Tuple[Population["50"], Individual]:
-    f1 = simulator.random_crosses(germplasm, 100)
+    f1, _ = simulator.random_crosses(germplasm, 100)
     dh_lines = simulator.double_haploid(f1, n_offspring=100)
     headrows = simulator.select(
         dh_lines, k=5, f_index=visual_selection(simulator, seed=7)

test.py

@@ -0,0 +1,6 @@
+from chromax import Simulator, sample_data
+simulator = Simulator(genetic_map=sample_data.genetic_map)
+f1 = simulator.load_population(sample_data.genome)
+f2, parent_ids = simulator.random_crosses(f1, 100, n_offspring=10)
+
+print(parent_ids.shape)

tests/test_simulator.py

@@ -148,11 +148,11 @@ def test_random_crosses():
     population = simulator.load_population(n_ind, ploidy=ploidy)
 
     n_crosses = 300
-    new_pop = simulator.random_crosses(population, n_crosses=n_crosses)
+    new_pop, _ = simulator.random_crosses(population, n_crosses=n_crosses)
     assert new_pop.shape == (n_crosses, n_markers, ploidy)
 
     n_offspring = 10
-    new_pop = simulator.random_crosses(
+    new_pop, _ = simulator.random_crosses(
         population=population, n_crosses=n_crosses, n_offspring=n_offspring
     )
     assert new_pop.shape == (n_crosses, n_offspring, n_markers, ploidy)
@@ -212,8 +212,8 @@ def test_seed_deterministic():
     mock_simulator = MockSimulator(n_markers=simulator1.n_markers)
     population = mock_simulator.load_population(n_ind, ploidy=ploidy)
 
-    new_pop1 = simulator1.random_crosses(population, n_crosses=10)
-    new_pop2 = simulator2.random_crosses(population, n_crosses=10)
+    new_pop1, _ = simulator1.random_crosses(population, n_crosses=10)
+    new_pop2, _ = simulator2.random_crosses(population, n_crosses=10)
 
     assert np.all(new_pop1 == new_pop2)