Update jungfrau notebooks

xfel_calibrate/calibrate.py

@@ -237,6 +237,46 @@ def get_notebook_function(nb, fname):
     return None
 
 
+def balance_sequences(in_folder, run, sequences, sequences_per_node,
+                      path_inset, max_nodes=8):
+    """
+    Return balance list of sequences to be executed on slurm nodes
+    Total list of sequences is splitted to several nodes, with a number of
+    sequences per node defined in the input parameter. if number
+    of required nodes is more than max_nodes, the number of sequences_per_node
+    will be increased to run on all on max_nodes.
+
+    :param in_folder: Path to raw data
+    :param run: Run number
+    :param sequences: List of sequences
+    :param sequences_per_node: Number of sequences per node
+    :param path_inset: Inset of path to data file
+    :param max_nodes: Maximum number of Maxwell nodes to use
+    :return: Balanced list of list of sequences
+    """
+    import glob
+    import re
+    import numpy as np
+    if sequences[0] == -1:
+        path = "{}/r{:04d}/*{}-S*.h5".format(in_folder, run, path_inset)
+        sequence_files = glob.glob(path)
+        seq_nums = set()
+        for sf in sequence_files:
+            seqnum = re.findall(r".*-S([0-9]*).h5", sf)[0]
+            seq_nums.add(int(seqnum))
+        seq_nums -= set(sequences)
+    else:
+        seq_nums = set(sequences)
+    nsplits = len(seq_nums) // sequences_per_node + 1
+    while nsplits > max_nodes:
+        sequences_per_node += 1
+        nsplits = len(seq_nums) // sequences_per_node + 1
+        print("Changed to {} sequences per node".format(sequences_per_node))
+        print("to have a maximum of 8 concurrent jobs")
+    return [l.tolist() for l in np.array_split(list(seq_nums), nsplits) if
+            l.size > 0]
+
+
 # extend the parser according to user input
 # the first case is if a detector was given, but no calibration type
 if len(sys.argv) == 3 and "-h" in sys.argv[2]:

xfel_calibrate/notebooks.py

@@ -162,7 +162,7 @@ notebooks = {
                                "concurrency": {"parameter": "sequences",
                                                "default concurrency": [-1],
                                                "use function": "balance_sequences",
-                                               "cluster cores": 4},
+                                               "cluster cores": 14},
                                 },
 
                        "STATS_FROM_DB":   {