From 6e0112a18cb4f157103535d6e963b77e2b295c34 Mon Sep 17 00:00:00 2001
From: samueldmcdermott <samueldmcdermott@gmail.com>
Date: Tue, 6 May 2025 15:08:00 -0400
Subject: [PATCH] setting seed

---
 deepszsim/dm_halo_dist.py |  7 ++++---
 deepszsim/noise.py        |  5 +++--
 deepszsim/simclusters.py  | 28 ++++++++++++++++++++--------
 deepszsim/simtools.py     | 11 ++++++++---
 4 files changed, 35 insertions(+), 16 deletions(-)

diff --git a/deepszsim/dm_halo_dist.py b/deepszsim/dm_halo_dist.py
index be72628..1f91413 100644
--- a/deepszsim/dm_halo_dist.py
+++ b/deepszsim/dm_halo_dist.py
@@ -33,9 +33,10 @@ def flatdist_halo(zmin, zmax, m200min_SM, m200max_SM, size, seed=None):
         distribution of random uniform redshifts starting at `m500min_SM` ending
         at `m500max_SM` with size `size`
     '''
-    _rng = np.random.default_rng(seed=seed)
+    rng = np.random.default_rng(seed=seed)
+    rng_M, rng_z = rng.spawn(2)
 
-    zdist=_rng.uniform(low=zmin, high=zmax, size=size)
-    mdist=_rng.uniform(low=m200min_SM, high=m200max_SM, size=size)
+    zdist=rng_z.uniform(low=zmin, high=zmax, size=size)
+    mdist=rng_M.uniform(low=m200min_SM, high=m200max_SM, size=size)
     
     return zdist, mdist
diff --git a/deepszsim/noise.py b/deepszsim/noise.py
index 60bbbfa..9e2e9dd 100644
--- a/deepszsim/noise.py
+++ b/deepszsim/noise.py
@@ -4,7 +4,7 @@
 
 import numpy as np
 
-def generate_noise_map(image_size, noise_level, pix_size):
+def generate_noise_map(image_size, noise_level, pix_size, seed):
     """
     Generates a white noise map based on the noise level and beam size.
 
@@ -21,7 +21,8 @@ def generate_noise_map(image_size, noise_level, pix_size):
     """
     
     # Create random noise map
-    random_noise_map = np.random.normal(0, 1, (image_size, image_size))
+    rng = np.random.default_rng(seed)
+    random_noise_map = rng.normal(0, 1, (image_size, image_size))
 
     # Scale random noise map by noise level
     scaled_noise_map = random_noise_map * noise_level
diff --git a/deepszsim/simclusters.py b/deepszsim/simclusters.py
index d945798..90573a6 100644
--- a/deepszsim/simclusters.py
+++ b/deepszsim/simclusters.py
@@ -54,6 +54,9 @@ def __init__(self, M200 = None, redshift_z = None, num_halos = None, halo_params
                  seed = None, tqverb = False
                  ):
         
+        self.rng_ = np.random.default_rng(seed)
+        self.rng, self.rng_ = self.rng_.spawn(2)
+        
         if (M200 is not None) and (redshift_z is not None):
             self.M200, self.redshift_z = M200, redshift_z
         else:
@@ -67,11 +70,13 @@ def __init__(self, M200 = None, redshift_z = None, num_halos = None, halo_params
             print(f"making {num_halos} clusters uniformly sampled from"
                   f"{halo_params_dict['zmin']}<z<{halo_params_dict['zmax']}, and "
                   f"{halo_params_dict['m200min_SM']:1.1e}<M200<{halo_params_dict['m200max_SM']:1.1e}")
+            randint = self.rng.integers(2**32-1)
+            self.rng, self.rng_ = self.rng_.spawn(2)
             self.redshift_z, self.M200 = dm_halo_dist.flatdist_halo(halo_params_dict['zmin'],
                                                                     halo_params_dict['zmax'],
                                                                     halo_params_dict['m200min_SM'],
                                                                     halo_params_dict['m200max_SM'],
-                                                                    int(num_halos), seed = seed)
+                                                                    int(num_halos), seed = randint)
         
         try:
             self._size = len(self.M200)
@@ -94,7 +99,7 @@ def __init__(self, M200 = None, redshift_z = None, num_halos = None, halo_params
         self.tqverb = tqverb
         
         self.alpha, self.gamma = alpha, gamma
-        self.seed, self._rng = seed, np.random.default_rng(seed)
+        self.rng, self.rng_ = self.rng_.spawn(2)
         
         if R200_Mpc is not None:
             self.R200_Mpc = R200_Mpc
@@ -104,7 +109,7 @@ def __init__(self, M200 = None, redshift_z = None, num_halos = None, halo_params
                                                            angsize_density = '500c')[1:3]
                  for i in range(self._size)]).T
         
-        self.id_list = [f"{int(self.M200[i]/1e9):07}_{int(self.redshift_z[i]*100):03}_{self._rng.integers(10**6):06}"
+        self.id_list = [f"{int(self.M200[i]/1e9):07}_{int(self.redshift_z[i]*100):03}_{self.rng.integers(10**6):06}"
                         for i in range(self._size)]
         self.clusters.update(zip(self.id_list, [{"params": {'M200': self.M200[i],
                                                             'redshift_z': self.redshift_z[i],
@@ -167,22 +172,29 @@ def get_T_maps(self, add_CMB = True, returnval = False):
         if add_CMB: self.ps = simtools.get_cls(ns = self.vars['ns'], cosmo = self.vars['cosmo'])
         if self.tqverb: print("making convolved T maps" + (" with CMB" if add_CMB else ""))
         _centerpix = self.image_size_pixels // 2
+        self.rng, self.rng_ = self.rng_.spawn(2)
+        seeds = self.rng.integers(2**32-1, size = self._size)
         for i in tqdm(range(self._size), disable = (not self.tqverb)):
             dTm, name = dT_maps[i], self.id_list[i]
             curdic = self.clusters[name]
             curdic['params']['dT_central'] = dTm[_centerpix,_centerpix]
             curdic['maps'] = {}
             beamsig_map = simtools.convolve_map_with_gaussian_beam(self.pixel_size_arcmin,
-                                                                   self.beam_size_arcmin, dTm)
+                                                                   self.beam_size_arcmin,
+                                                                   dTm)
             if add_CMB:
-                conv_map, cmb_map = simtools.add_cmb_map_and_convolve(dTm, self.ps,
+                conv_map, cmb_map = simtools.add_cmb_map_and_convolve(dTm,
+                                                                      self.ps,
                                                                       self.pixel_size_arcmin,
-                                                                      self.beam_size_arcmin)
+                                                                      self.beam_size_arcmin,
+                                                                      seed = int(seeds[i]))
             else:
                 conv_map, cmb_map = beamsig_map, np.zeros_like(beamsig_map)
             if not self.vars['noise_level'] == 0:
-                noise_map = noise.generate_noise_map(self.image_size_pixels, self.vars['noise_level'],
-                                                     self.pixel_size_arcmin)
+                noise_map = noise.generate_noise_map(self.image_size_pixels,
+                                                     self.vars['noise_level'],
+                                                     self.pixel_size_arcmin,
+                                                     seed = int(seeds[i]))
             else:
                 noise_map = np.zeros_like(conv_map)
             final_map = conv_map + noise_map
diff --git a/deepszsim/simtools.py b/deepszsim/simtools.py
index c5203fe..336ef5b 100644
--- a/deepszsim/simtools.py
+++ b/deepszsim/simtools.py
@@ -59,8 +59,11 @@ def f_sz(freq_ghz, T_CMB_K):
 
     return fsz
 
-def add_cmb_map_and_convolve(dT_map_uK, ps, pix_size_arcmin, 
-                                 beam_size_fwhp_arcmin):
+def add_cmb_map_and_convolve(dT_map_uK,
+                             ps,
+                             pix_size_arcmin,
+                             beam_size_fwhp_arcmin,
+                             seed):
     '''
     Parameters:
     ----------
@@ -83,7 +86,9 @@ def add_cmb_map_and_convolve(dT_map_uK, ps, pix_size_arcmin,
                         dT_map_uK.shape[1]+2*padding_value)
     #print(expanded_shape)
     cmb_map = make_cmb_map(shape=expanded_shape, 
-                                pix_size_arcmin=pix_size_arcmin, ps=ps)
+                           pix_size_arcmin=pix_size_arcmin,
+                           ps=ps,
+                           seed = seed)
     if type(dT_map_uK) is u.Quantity:
         cmb_map = cmb_map *u.uK
     dT_map_expanded = np.pad(dT_map_uK, (padding_value,padding_value),