new function for random sampling

pySD/sdmout_src/sdtracing.py

@@ -6,7 +6,7 @@
 Author: Clara Bayley (CB)
 Additional Contributors:
 -----
-Last Modified: Monday 20th November 2023
+Last Modified: Tuesday 16th January 2024
 Modified By: CB
 -----
 License: BSD 3-Clause "New" or "Revised" License
@@ -59,7 +59,7 @@ def attribute_for_superdroplets_sample(sddata, attr, ndrops2sample=0,
   ''' returns 2D array with dimensions [time, SD]
   containing attribute data over time for a sample of 
   superdroplets. Sample is either for superdroplets with
-  specific Ids in 'ids' list, or sample if 'ndrops2sample'
+  specific Ids in 'ids' list, or sample of 'ndrops2sample'
   randomly selected superdrops with Ids in the range 
   [minid, maxid] '''
 
@@ -99,4 +99,27 @@ def attributes_at_times(sddata, time, times2sel, attrs2sel):
     selattr_data = attr_at_times(sddata[attr], time, times2sel)
     selected_data[attr] = selattr_data
 
-  return selected_data
+  return selected_data
+
+def attrs_for_superdroplets_sample(sddata, attrs, ndrops2sample=0,
+                                    minid=0, maxid=0, ids=[]):
+  ''' returns dictionary of 2D arrays (with dimensions [time, SD])
+  for each attribute in 'attrs' list for a sample of 
+  superdroplets. Sample is either for superdroplets with
+  specific Ids in 'ids' list, or sample of 'ndrops2sample'
+  randomly selected superdrops with Ids in the range 
+  [minid, maxid] '''
+
+  if ids == []:
+    population = list(range(minid, maxid, 1))
+    if ndrops2sample == 0:
+      ndrops2sample = maxid
+    sample = random.sample(population, ndrops2sample)
+  else:
+    sample = ids
+
+  data = {}
+  for a in attrs: 
+    data[a] = attribute_for_superdroplets_sample(sddata, a, ids=sample)
+
+  return data