added support for channel signals (i.e. signals that come as channels but we only want individual chords

Julian Kates-Harbeck · Julian Kates-Harbeck · commit 14858b92a3e9 · 2017-12-14T21:20:48.000-05:00
diff --git a/data/signals.py b/data/signals.py
@@ -3,7 +3,7 @@
 import time
 import sys
 
-from plasma.primitives.data import Signal,ProfileSignal,Machine
+from plasma.primitives.data import Signal,ProfileSignal,ChannelSignal,Machine
 
 def create_missing_value_filler():
 	time = np.linspace(0,100,100)
@@ -176,8 +176,8 @@ def fetch_nstx_data(signal_path,shot_num,c):
 pin = Signal("Input Power (beam for d3d)",['jpf/gs/bl-ptot<s','d3d/bmspinj'],[jet,d3d]) #Total Beam Power
 
 pradtot = Signal("Radiated Power",['jpf/db/b5r-ptot>out'],[jet])
-pradcore = Signal("Radiated Power Core",['d3d/'+r'\bol_l15_p'],[d3d])
-pradedge = Signal("Radiated Power Edge",['d3d/'+r'\bol_l03_p'],[d3d])
+pradcore = ChannelSignal("Radiated Power Core",['ppf/bolo/kb5h/channel14', 'd3d/'+r'\bol_l15_p'],[jet,d3d])
+pradedge = ChannelSignal("Radiated Power Edge",['ppf/bolo/kb5h/channel10','d3d/'+r'\bol_l03_p'],[jet,d3d])
 # pechin = Signal("ECH input power, not always on",['d3d/pcechpwrf'],[d3d])
 pechin = Signal("ECH input power, not always on",['RF/ECH.TOTAL.ECHPWRC'],[d3d])
 
diff --git a/plasma/primitives/data.py b/plasma/primitives/data.py
@@ -118,6 +118,27 @@ def load_data(self,prepath,shot,dtype='float32'):
 
         return t,sig,True
 
+
+    def fetch_data_basic(self,machine,shot_num,c):
+        path = self.get_path(machine)
+        mapping_path = self.get_mapping_path(machine)
+        success = False
+        mapping = None
+        try:
+            time,data,mapping,success = machine.fetch_data_fn(path,shot_num,c)
+        except Exception as e:
+            print(e)
+            sys.stdout.flush()
+
+        if not success:
+            return None,None,None,False
+
+        time = np.array(time) + 1e-3*self.get_causal_shift(machine)
+        return time,np.array(data),mapping,success    
+
+    def fetch_data(self,machine,shot_num,c):
+        return self.fetch_data_basic(machine,shot_num,c)
+
     def is_defined_on_machine(self,machine):
         return machine in self.machines
 
@@ -206,6 +227,75 @@ def load_data(self,prepath,shot,dtype='float32'):
 
         return t,sig_interp,True
 
+    def fetch_data(self,machine,shot_num,c):
+        time,data,mapping,success = self.fetch_data_basic(machine,shot_num,c)
+
+        if mapping is not None and np.ndim(mapping) == 1:#make sure there is a mapping for every timestep
+            T = len(time)
+            mapping = np.tile(mapping,(T,1)).transpose()
+            assert(mapping.shape == data.shape), "shape of mapping and data is different"
+        if mapping_path is not None:#fetch the mapping separately
+            time_map,data_map,mapping_map,success_map = machine.fetch_data_fn(mapping_path,shot_num,c)
+            success = (success and success_map)
+            if not success:
+                print("No success for signal {} and mapping {}".format(path,mapping_path))
+            else:
+                assert(np.all(time == time_map)), "time for signal {} and mapping {} don't align: \n{}\n\n{}\n".format(path,mapping_path,time,time_map)
+                mapping = data_map
+
+        if not success:
+            return None,None,None,False
+        return time,data,mapping,success 
+
+
+class ChannelSignal(Signal):
+    def __init__(self,description,paths,machines,tex_label=None,causal_shifts=None,mapping_range=(0,1),num_channels=32,data_avail_tolerances=None,is_strictly_positive=False,mapping_paths=None):
+        super(ChannelSignal, self).__init__(description,paths,machines,tex_label,causal_shifts,is_ip=False,data_avail_tolerances=data_avail_tolerances,is_strictly_positive=is_strictly_positive,mapping_paths=mapping_paths)
+        nums,new_paths = self.get_channel_nums(paths)
+        self.channel_nums = nums
+        self.paths = new_paths
+
+    def get_channel_nums(self,paths):
+        regex = re.compile('channel\d+')
+        regex_int = re.compile('\d+')
+        nums = []
+        new_paths = []
+        for p in paths:
+            assert(p[-1] != '/')
+            elements = p.split('/')
+            res = regex.findall(elements[-1])
+            assert(len(res) < 2)
+            if len(res) == 0:
+                nums.append(None)
+                new_paths.append(p)
+            else:
+                nums.append(int(regex_int.findall(res[0])[0]))
+                new_paths.append("/".join(elements[:-1]))
+        return nums,new_paths
+
+    def get_channel_num(self,machine):
+        idx = self.get_idx(machine)
+        return self.channel_nums[idx]
+
+    def fetch_data(self,machine,shot_num,c):
+        time,data,mapping,success = self.fetch_data_basic(machine,shot_num,c)
+        channel_num = self.get_channel_num(machine)
+        if channel_num is not None and success:
+            if np.ndim(data) != 2:
+                print("Channel Signal {} expected 2D array for shot {}".format(self,shot))
+                success = False
+            else:
+                data = data[:,channel_num]
+        return time,data,mapping,success
+
+    def get_file_path(self,prepath,machine,shot_number):
+        dirname = self.get_path(machine)
+        num = self.get_channel_num(machine)
+        if num is not None:
+            dirname += "/channel{}".format(num)
+        return get_individual_shot_file(prepath + '/' + machine.name + '/' +dirname + '/',shot_number)
+
+
 
 class Machine(object):
     def __init__(self,name,server,fetch_data_fn,max_cores = 8,current_threshold=0):
@@ -218,35 +308,6 @@ def __init__(self,name,server,fetch_data_fn,max_cores = 8,current_threshold=0):
     def get_connection(self):
         return Connection(server)
 
-    def fetch_data(self,signal,shot_num,c):
-        path = signal.get_path(self)
-        mapping_path = signal.get_mapping_path(self)
-        success = False
-        mapping = None
-        try:
-            time,data,mapping,success = self.fetch_data_fn(path,shot_num,c)
-            if mapping is not None and np.ndim(mapping) == 1:#make sure there is a mapping for every timestep
-                T = len(time)
-                mapping = np.tile(mapping,(T,1)).transpose()
-                assert(mapping.shape == data.shape), "shape of mapping and data is different"
-            if mapping_path is not None:#fetch the mapping separately
-                time_map,data_map,mapping_map,success_map = self.fetch_data_fn(mapping_path,shot_num,c)
-                success = (success and success_map)
-                if not success:
-                    print("No success for signal {} and mapping {}".format(path,mapping_path))
-                else:
-               	    assert(np.all(time == time_map)), "time for signal {} and mapping {} don't align: \n{}\n\n{}\n".format(path,mapping_path,time,time_map)
-                    mapping = data_map
-        except Exception as e:
-            print(e)
-            sys.stdout.flush()
-
-        if not success:
-            return None,None,None,False
-
-        time = np.array(time) + 1e-3*signal.get_causal_shift(self)
-        return time,np.array(data),mapping,success
-
     def __eq__(self,other):
         return self.name.__eq__(other.name)
 
diff --git a/plasma/utils/downloading.py b/plasma/utils/downloading.py
@@ -90,7 +90,7 @@ def save_shot(shot_num_queue,c,signals,save_prepath,machine,sentinel=-1):
             if not success:
                 try:
                     try:
-                        time,data,mapping,success = machine.fetch_data(signal,shot_num,c) 
+                        time,data,mapping,success = signal.fetch_data(machine,shot_num,c) 
                         if not success:
                             print('No success shot {}, signal {}'.format(shot_num,signal))
                     except:
