diff --git a/software/grbl/grbl_interactive.py b/software/grbl/grbl_interactive.py index 83c385d9..acbc4552 100644 --- a/software/grbl/grbl_interactive.py +++ b/software/grbl/grbl_interactive.py @@ -9,214 +9,237 @@ class GRBLManager: # a is the left motor on the wall # b is the right motor on the wall # origin is very close to b - #origin_a=np.array([2880,-300]) # a1 / y - #origin_b=np.array([-300,-365]) # a2 / x - #origin_a=np.array([2910,-350]) # a1 / y - origin_a=np.array([2910+100,-350]) # a1 / y - origin_b=np.array([-300,-370]) # a2 / x + # origin_a=np.array([2880,-300]) # a1 / y + # origin_b=np.array([-300,-365]) # a2 / x + # origin_a=np.array([2910,-350]) # a1 / y + origin_a = np.array([2910 + 100, -350]) # a1 / y + origin_b = np.array([-300, -370]) # a2 / x - limit_topleft=np.array([2880-300,4000-200]) - limit_bottomright=np.array([0,1500-200]) - x_range=2880-300 - y_max=2500 + limit_topleft = np.array([2880 - 300, 4000 - 200]) + limit_bottomright = np.array([0, 1500 - 200]) + x_range = 2880 - 300 + y_max = 2500 - def from_steps(self,a_motor_steps,b_motor_steps): - r1=np.linalg.norm(self.origin_a)-a_motor_steps - r2=np.linalg.norm(self.origin_b)-b_motor_steps - d=np.linalg.norm(self.origin_a-self.origin_b) + def from_steps(self, a_motor_steps, b_motor_steps): + r1 = np.linalg.norm(self.origin_a) - a_motor_steps + r2 = np.linalg.norm(self.origin_b) - b_motor_steps + d = np.linalg.norm(self.origin_a - self.origin_b) - x=(d*d-r2*r2+r1*r1)/(2*d) # on different axis - x_dir=(self.origin_b-self.origin_a)/d + x = (d * d - r2 * r2 + r1 * r1) / (2 * d) # on different axis + x_dir = (self.origin_b - self.origin_a) / d - y=np.sqrt(r1*r1-x*x) - y_dir=np.array([-x_dir[1],x_dir[0]]) #x_dir.T # orthogonal to x + y = np.sqrt(r1 * r1 - x * x) + y_dir = np.array([-x_dir[1], x_dir[0]]) # x_dir.T # orthogonal to x - xy_g=self.origin_a-y_dir*y+x_dir*x + xy_g = self.origin_a - y_dir * y + x_dir * x return xy_g - def to_steps(self,p): - x_frac=p[0]/self.x_range - y_limit=min(self.y_max,(self.limit_bottomright*(1-x_frac)+self.limit_topleft*x_frac)[1]) - - p[0]=min(self.x_range,max(0,p[0])) - p[1]=min(y_limit,max(0,p[1])) + def to_steps(self, p): + x_frac = p[0] / self.x_range + y_limit = min( + self.y_max, + (self.limit_bottomright * (1 - x_frac) + self.limit_topleft * x_frac)[1], + ) - a_motor_steps=np.linalg.norm(self.origin_a)-np.linalg.norm(self.origin_a-p) + p[0] = min(self.x_range, max(0, p[0])) + p[1] = min(y_limit, max(0, p[1])) - b_motor_steps=np.linalg.norm(self.origin_b)-np.linalg.norm(self.origin_b-p)#-np.linalg.norm(a2) - return a_motor_steps,b_motor_steps + a_motor_steps = np.linalg.norm(self.origin_a) - np.linalg.norm( + self.origin_a - p + ) + + b_motor_steps = np.linalg.norm(self.origin_b) - np.linalg.norm( + self.origin_b - p + ) # -np.linalg.norm(a2) + return a_motor_steps, b_motor_steps def spiral(self): - center=np.array([1500,900]) - spiral_radius=1200 - t_max=6*2*np.pi - v=spiral_radius/t_max - w=1 - for t in np.linspace(0,t_max,256*16*2): - x=(v*t)*np.cos(w*t) - y=(v*t)*np.sin(w*t) - p=np.array([x,y])+center - a_motor_steps,b_motor_steps=to_steps(p) - cmd="G0 X%0.2f Y%0.2f" % (b_motor_steps,a_motor_steps) - #print("SENDING",x,y,cmd) - s.write((cmd + '\n').encode()) # Send g-code block to grbl + center = np.array([1500, 900]) + spiral_radius = 1200 + t_max = 6 * 2 * np.pi + v = spiral_radius / t_max + w = 1 + for t in np.linspace(0, t_max, 256 * 16 * 2): + x = (v * t) * np.cos(w * t) + y = (v * t) * np.sin(w * t) + p = np.array([x, y]) + center + a_motor_steps, b_motor_steps = to_steps(p) + cmd = "G0 X%0.2f Y%0.2f" % (b_motor_steps, a_motor_steps) + # print("SENDING",x,y,cmd) + s.write((cmd + "\n").encode()) # Send g-code block to grbl s.readline().strip() def calibrate(self): - for x in np.linspace(1500,500,5): - for y in np.linspace(30,1200,5): - x,y=to_steps(np.array([x,y])) - cmd="G0 X%0.2f Y%0.2f" % (x,y) - s.write((cmd + '\n').encode()) # Send g-code block to grbl + for x in np.linspace(1500, 500, 5): + for y in np.linspace(30, 1200, 5): + x, y = to_steps(np.array([x, y])) + cmd = "G0 X%0.2f Y%0.2f" % (x, y) + s.write((cmd + "\n").encode()) # Send g-code block to grbl - def __init__(self,serial_fn): + def __init__(self, serial_fn): # Open grbl serial port ==> CHANGE THIS BELOW TO MATCH YOUR USB LOCATION - self.s = serial.Serial(serial_fn,115200,timeout=0.3,write_timeout=0.3) # GRBL operates at 115200 baud. Leave that part alone. + self.s = serial.Serial( + serial_fn, 115200, timeout=0.3, write_timeout=0.3 + ) # GRBL operates at 115200 baud. Leave that part alone. self.s.write("?".encode()) - grbl_out=self.s.readline() # get the response - print("GRBL ONLINE",grbl_out) - self.position={'time':time.time(),'xy':np.zeros(2)} + grbl_out = self.s.readline() # get the response + print("GRBL ONLINE", grbl_out) + self.position = {"time": time.time(), "xy": np.zeros(2)} self.update_status() time.sleep(0.05) - self.collect=True + self.collect = True def push_reset(self): - self.s.write(b'\x18') + self.s.write(b"\x18") return self.s.readline().decode().strip() - def update_status(self,skip_write=False): + def update_status(self, skip_write=False): if not skip_write: time.sleep(0.01) self.s.write("?".encode()) time.sleep(0.01) - start_time=time.time() - response=self.s.readline().decode().strip() + start_time = time.time() + response = self.s.readline().decode().strip() time.sleep(0.01) - #print("STATUS",response) - # + # print("STATUS",response) + # try: - motor_position_str=response.split("|")[1] + motor_position_str = response.split("|")[1] except Exception as e: - print("FAILED TO PARSE",response,"|e|",e,time.time()-start_time) + print("FAILED TO PARSE", response, "|e|", e, time.time() - start_time) return self.update_status(skip_write=not skip_write) - b_motor_steps,a_motor_steps,_,_=map(float,motor_position_str[len('MPos:'):].split(',')) + b_motor_steps, a_motor_steps, _, _ = map( + float, motor_position_str[len("MPos:") :].split(",") + ) - xy=self.from_steps(a_motor_steps,b_motor_steps) - is_moving=(self.position['xy']!=xy).any() - self.position={ - 'a_motor_steps':a_motor_steps, - 'b_motor_steps':b_motor_steps, - 'xy':xy, - 'is_moving':is_moving, - 'time':time.time() - } + xy = self.from_steps(a_motor_steps, b_motor_steps) + is_moving = (self.position["xy"] != xy).any() + self.position = { + "a_motor_steps": a_motor_steps, + "b_motor_steps": b_motor_steps, + "xy": xy, + "is_moving": is_moving, + "time": time.time(), + } return self.position def wait_while_moving(self): while True: - old_pos=self.update_status() + old_pos = self.update_status() time.sleep(0.05) - new_pos=self.update_status() - if old_pos['a_motor_steps']==new_pos['a_motor_steps'] and old_pos['b_motor_steps']==new_pos['b_motor_steps']: - return + new_pos = self.update_status() + if ( + old_pos["a_motor_steps"] == new_pos["a_motor_steps"] + and old_pos["b_motor_steps"] == new_pos["b_motor_steps"] + ): + return time.sleep(0.01) def binary_search_edge(self, left, right, xy, direction, epsilon): - if (right-left)0.0001: - #go back left - return self.binary_search_edge(left,l,xy, direction, epsilon) - return self.binary_search_edge(l,right,xy,direction,epsilon) + l = (right + left) / 2 + p = l * direction + xy + steps = self.to_steps(np.copy(p)) + actual = self.from_steps(*steps) + deviation = np.linalg.norm(p - actual) + if deviation > 0.0001: + # go back left + return self.binary_search_edge(left, l, xy, direction, epsilon) + return self.binary_search_edge(l, right, xy, direction, epsilon) - def bounce(self,bounces,direction=None): + def bounce(self, bounces, direction=None): if direction is None: - theta=np.random.uniform(2*np.pi) - direction=np.array([np.sin(theta),np.cos(theta)]) - percent_random=0.05 - theta=np.random.uniform(2*np.pi) - direction=(1-percent_random)*direction+percent_random*np.array([np.sin(theta),np.cos(theta)]) - direction=direction/np.linalg.norm(direction) - print("bounce",direction,np.linalg.norm(direction)) + theta = np.random.uniform(2 * np.pi) + direction = np.array([np.sin(theta), np.cos(theta)]) + percent_random = 0.05 + theta = np.random.uniform(2 * np.pi) + direction = (1 - percent_random) * direction + percent_random * np.array( + [np.sin(theta), np.cos(theta)] + ) + direction = direction / np.linalg.norm(direction) + print("bounce", direction, np.linalg.norm(direction)) for _ in range(bounces): - to_points,new_direction=self.single_bounce(direction) - if len(to_points)>0: - print("MOVE",to_points[0],to_points[-1]) + to_points, new_direction = self.single_bounce(direction) + if len(to_points) > 0: + print("MOVE", to_points[0], to_points[-1]) for point in to_points: self.move_to(point) - #print("MOVE") + # print("MOVE") self.update_status() - while np.linalg.norm(self.position['xy']-point)>100: + while np.linalg.norm(self.position["xy"] - point) > 100: self.update_status() - if (new_direction!=direction).any(): # we are changing direction + if (new_direction != direction).any(): # we are changing direction self.wait_while_moving() - direction=new_direction + direction = new_direction return direction - def single_bounce(self,direction,xy=None,step_size=5): - #find current position - #pick a random direction - #take full field step - #if hit a wall - #direction=np.array([1,0]) - #direction=direction/np.linalg.norm(direction) + def single_bounce(self, direction, xy=None, step_size=5): + # find current position + # pick a random direction + # take full field step + # if hit a wall + # direction=np.array([1,0]) + # direction=direction/np.linalg.norm(direction) if xy is None: self.update_status() - xy=self.position['xy'] - #find out at what point xy+l*direction we stop changing one of the variables - epsilon=0.001 - l=self.binary_search_edge(0,10000,xy,direction,epsilon) - #find a paralell vector to the boundary - p1=self.from_steps(*self.to_steps((l+2*epsilon)*direction+xy)) - p2=self.from_steps(*self.to_steps((l+3*epsilon)*direction+xy)) - if np.linalg.norm(p1-p2) CHANGE THIS BELOW TO MATCH YOUR USB LOCATION -s = serial.Serial(serial_fn,115200) # GRBL operates at 115200 baud. Leave that part alone. +s = serial.Serial( + serial_fn, 115200 +) # GRBL operates at 115200 baud. Leave that part alone. # Open g-code file -f = open(gcode_fn,'r'); +f = open(gcode_fn, "r") # Wake up grbl s.write("?".encode()) print(s.readline().strip()) -s.write(b'\x18') +s.write(b"\x18") print(s.readline().strip()) s.write("?".encode()) print(s.readline().strip()) s.write("?".encode()) print(s.readline().strip()) -#s.write("G0 X100\n".encode()) -#print(s.readline().strip()) +# s.write("G0 X100\n".encode()) +# print(s.readline().strip()) -#sys.exit(1) -#time.sleep(2) # Wait for grbl to initialize -#s.flushInput() # Flush startup text in serial input +# sys.exit(1) +# time.sleep(2) # Wait for grbl to initialize +# s.flushInput() # Flush startup text in serial input # Stream g-code to grbl for line in f: - l = line.strip() # Strip all EOL characters for consistency - if l[0]==';': - continue - print('Sending: ' + l) - s.write((l + '\n').encode()) # Send g-code block to grbl - grbl_out = s.readline() # Wait for grbl response with carriage return - print(grbl_out) - - # Wait here until grbl is finished to close serial port and file. - #raw_input(" Press to exit and disable grbl.") + l = line.strip() # Strip all EOL characters for consistency + if l[0] == ";": + continue + print("Sending: " + l) + s.write((l + "\n").encode()) # Send g-code block to grbl + grbl_out = s.readline() # Wait for grbl response with carriage return + print(grbl_out) + + # Wait here until grbl is finished to close serial port and file. + # raw_input(" Press to exit and disable grbl.") # Close file and serial port f.close() diff --git a/software/grbl_sdr_collect.py b/software/grbl_sdr_collect.py index ea5e0b92..699dcec7 100644 --- a/software/grbl_sdr_collect.py +++ b/software/grbl_sdr_collect.py @@ -1,5 +1,5 @@ -from sdrpluto.gather import * -from grbl.grbl_interactive import GRBLManager +from sdrpluto.gather import * +from grbl.grbl_interactive import GRBLManager from model_training_and_inference.utils.rf import beamformer import threading import time @@ -8,106 +8,141 @@ import os import pickle + def bounce_grbl(gm): - direction=None + direction = None while gm.collect: print("TRY TO BOUNCE") try: - direction=gm.bounce(1,direction=direction) + direction = gm.bounce(1, direction=direction) except Exception as e: print(e) print("TRY TO BOUNCE RET") - time.sleep(10) # cool off the motor + time.sleep(10) # cool off the motor + -if __name__=='__main__': +if __name__ == "__main__": parser = argparse.ArgumentParser() - parser.add_argument("--receiver-ip", type=str, help="receiver Pluto IP address",required=True) - parser.add_argument("--emitter-ip", type=str, help="emitter Pluto IP address",required=True) - parser.add_argument("--fi", type=int, help="Intermediate frequency",required=False,default=1e5) - parser.add_argument("--fc", type=int, help="Carrier frequency",required=False,default=2.5e9) - parser.add_argument("--fs", type=int, help="Sampling frequency",required=False,default=16e6) - parser.add_argument("--cal0", type=int, help="Rx0 calibration phase offset in degrees",required=False,default=180) - #parser.add_argument("--d", type=int, help="Distance apart",required=False,default=0.062) - parser.add_argument("--rx-gain", type=int, help="RX gain",required=False,default=-3) - parser.add_argument("--tx-gain", type=int, help="TX gain",required=False,default=-3) - parser.add_argument("--grbl-serial", type=str, help="serial file for GRBL",required=True) - parser.add_argument("--out", type=str, help="output file",required=True) - parser.add_argument("--record-freq", type=int, help="record freq",required=False,default=5) - parser.add_argument("--rx-mode", type=str, help="rx mode",required=False,default="fast_attack") - parser.add_argument("--record-n", type=int, help="records",required=False,default=43200) - parser.add_argument("--rx-n", type=int, help="RX buffer size",required=False,default=2**12) + parser.add_argument( + "--receiver-ip", type=str, help="receiver Pluto IP address", required=True + ) + parser.add_argument( + "--emitter-ip", type=str, help="emitter Pluto IP address", required=True + ) + parser.add_argument( + "--fi", type=int, help="Intermediate frequency", required=False, default=1e5 + ) + parser.add_argument( + "--fc", type=int, help="Carrier frequency", required=False, default=2.5e9 + ) + parser.add_argument( + "--fs", type=int, help="Sampling frequency", required=False, default=16e6 + ) + parser.add_argument( + "--cal0", + type=int, + help="Rx0 calibration phase offset in degrees", + required=False, + default=180, + ) + # parser.add_argument("--d", type=int, help="Distance apart",required=False,default=0.062) + parser.add_argument( + "--rx-gain", type=int, help="RX gain", required=False, default=-3 + ) + parser.add_argument( + "--tx-gain", type=int, help="TX gain", required=False, default=-3 + ) + parser.add_argument( + "--grbl-serial", type=str, help="serial file for GRBL", required=True + ) + parser.add_argument("--out", type=str, help="output file", required=True) + parser.add_argument( + "--record-freq", type=int, help="record freq", required=False, default=5 + ) + parser.add_argument( + "--rx-mode", type=str, help="rx mode", required=False, default="fast_attack" + ) + parser.add_argument( + "--record-n", type=int, help="records", required=False, default=43200 + ) + parser.add_argument( + "--rx-n", type=int, help="RX buffer size", required=False, default=2**12 + ) args = parser.parse_args() - #setup output recorder - record_matrix = np.memmap(args.out, dtype='float32', mode='w+', shape=(args.record_n,5+65)) + # setup output recorder + record_matrix = np.memmap( + args.out, dtype="float32", mode="w+", shape=(args.record_n, 5 + 65) + ) - #setup GRBL - gm=GRBLManager(args.grbl_serial) + # setup GRBL + gm = GRBLManager(args.grbl_serial) - #calibrate the receiver - sdr_rx=setup_rxtx_and_phase_calibration(args) - if sdr_rx==None: + # calibrate the receiver + sdr_rx = setup_rxtx_and_phase_calibration(args) + if sdr_rx == None: print("Failed phase calibration, exiting") sys.exit(1) - phase_calibration=sdr_rx.phase_calibration - sdr_rx=None - sdr_rx,sdr_tx=setup_rx_and_tx(args) - if sdr_rx==None or sdr_tx==None: + phase_calibration = sdr_rx.phase_calibration + sdr_rx = None + sdr_rx, sdr_tx = setup_rx_and_tx(args) + if sdr_rx == None or sdr_tx == None: print("Failed setup, exiting") sys.exit(1) - #apply the previous calibration - sdr_rx.phase_calibration=phase_calibration + # apply the previous calibration + sdr_rx.phase_calibration = phase_calibration - #start moving GRBL + # start moving GRBL gm_thread = threading.Thread(target=bounce_grbl, args=(gm,)) gm_thread.start() - pos=np.array([[-0.03,0],[0.03,0]]) + pos = np.array([[-0.03, 0], [0.03, 0]]) - time_offset=time.time() + time_offset = time.time() for idx in range(args.record_n): - while not gm.position['is_moving']: + while not gm.position["is_moving"]: print("wait for movement") time.sleep(1) - #get some data + # get some data try: - signal_matrix=sdr_rx.rx() + signal_matrix = sdr_rx.rx() except Exception as e: - print("Failed to receive RX data! removing file",e) + print("Failed to receive RX data! removing file", e) os.remove(args.out) break - signal_matrix[1]*=np.exp(1j*sdr_rx.phase_calibration) - current_time=time.time()-time_offset # timestamp + signal_matrix[1] *= np.exp(1j * sdr_rx.phase_calibration) + current_time = time.time() - time_offset # timestamp - beam_thetas,beam_sds,beam_steer=beamformer( - pos, - signal_matrix, - args.fc - ) + beam_thetas, beam_sds, beam_steer = beamformer(pos, signal_matrix, args.fc) - avg_phase_diff=get_avg_phase(signal_matrix) - xy=gm.position['xy'] - record_matrix[idx]=np.hstack( - [ - np.array([current_time,xy[0],xy[1],avg_phase_diff[0],avg_phase_diff[1]]), - beam_sds]) - #time.sleep(1.0/args.record_freq) + avg_phase_diff = get_avg_phase(signal_matrix) + xy = gm.position["xy"] + record_matrix[idx] = np.hstack( + [ + np.array( + [current_time, xy[0], xy[1], avg_phase_diff[0], avg_phase_diff[1]] + ), + beam_sds, + ] + ) + # time.sleep(1.0/args.record_freq) - if idx%50==0: - elapsed=time.time()-time_offset - rate=elapsed/(idx+1) - print(idx, - "mean: %0.4f" % avg_phase_diff[0], - "_mean: %0.4f" % avg_phase_diff[1], - "%0.4f" % (100.0*float(idx)/args.record_n), - '%', - "elapsed(min): %0.1f" % (elapsed/60), - "rate(s/idx): %0.3f" % rate, - "remaining(min): %0.3f" % ((args.record_n-idx)*rate/60)) - gm.collect=False + if idx % 50 == 0: + elapsed = time.time() - time_offset + rate = elapsed / (idx + 1) + print( + idx, + "mean: %0.4f" % avg_phase_diff[0], + "_mean: %0.4f" % avg_phase_diff[1], + "%0.4f" % (100.0 * float(idx) / args.record_n), + "%", + "elapsed(min): %0.1f" % (elapsed / 60), + "rate(s/idx): %0.3f" % rate, + "remaining(min): %0.3f" % ((args.record_n - idx) * rate / 60), + ) + gm.collect = False print("Done collecting!") gm_thread.join() - #plot_recv_signal(sdr_rx) - + # plot_recv_signal(sdr_rx) diff --git a/software/model_training_and_inference/01_generate_data.py b/software/model_training_and_inference/01_generate_data.py index 3a630c5d..5a6b9dd8 100644 --- a/software/model_training_and_inference/01_generate_data.py +++ b/software/model_training_and_inference/01_generate_data.py @@ -8,55 +8,80 @@ from tqdm import tqdm from compress_pickle import dump, load -from utils.spf_generate import generate_session_and_dump,generate_session +from utils.spf_generate import generate_session_and_dump, generate_session -if __name__=='__main__': - parser = argparse.ArgumentParser() - parser.add_argument('--carrier-frequency', type=float, required=False, default=2.4e9) - parser.add_argument('--signal-frequency', type=float, required=False, default=100e3) - parser.add_argument('--sampling-frequency', type=float, required=False, default=10e6) - parser.add_argument('--array-type', type=str, required=False, default='circular',choices=['linear','circular']) - parser.add_argument('--elements', type=int, required=False, default=11) - parser.add_argument('--random-silence', type=bool, required=False, default=False) - parser.add_argument('--detector-noise', type=float, required=False, default=1e-4) - parser.add_argument('--random-emitter-timing', type=bool, required=False, default=False) - parser.add_argument('--sources', type=int, required=False, default=2) - parser.add_argument('--seed', type=int, required=False, default=0) - parser.add_argument('--numba', type=bool, required=False, default=False) - parser.add_argument('--beam-former-spacing', type=int, required=False, default=256+1) - parser.add_argument('--width', type=int, required=False, default=128) - parser.add_argument('--detector-trajectory', type=str, required=False, default='bounce',choices=['orbit','bounce']) - parser.add_argument('--detector-speed', type=float, required=False, default=10.0) - parser.add_argument('--source-speed', type=float, required=False, default=0.0) - parser.add_argument('--sigma', type=float, required=False, default=1.0) - parser.add_argument('--time-steps', type=int, required=False, default=100) - parser.add_argument('--time-interval', type=float, required=False, default=0.3) - parser.add_argument('--samples-per-snapshot', type=int, required=False, default=3) - parser.add_argument('--sessions', type=int, required=False, default=1024) - parser.add_argument('--output', type=str, required=False, default="sessions-default") - parser.add_argument('--reference', type=bool, required=False, default=False) - parser.add_argument('--cpus', type=int, required=False, default=8) - parser.add_argument('--live', type=bool, required=False, default=False) - parser.add_argument('--profile', type=bool, required=False, default=False) - parser.add_argument('--fixed-detector', type=float, required=False, nargs='+') +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--carrier-frequency", type=float, required=False, default=2.4e9 + ) + parser.add_argument("--signal-frequency", type=float, required=False, default=100e3) + parser.add_argument( + "--sampling-frequency", type=float, required=False, default=10e6 + ) + parser.add_argument( + "--array-type", + type=str, + required=False, + default="circular", + choices=["linear", "circular"], + ) + parser.add_argument("--elements", type=int, required=False, default=11) + parser.add_argument("--random-silence", type=bool, required=False, default=False) + parser.add_argument("--detector-noise", type=float, required=False, default=1e-4) + parser.add_argument( + "--random-emitter-timing", type=bool, required=False, default=False + ) + parser.add_argument("--sources", type=int, required=False, default=2) + parser.add_argument("--seed", type=int, required=False, default=0) + parser.add_argument("--numba", type=bool, required=False, default=False) + parser.add_argument( + "--beam-former-spacing", type=int, required=False, default=256 + 1 + ) + parser.add_argument("--width", type=int, required=False, default=128) + parser.add_argument( + "--detector-trajectory", + type=str, + required=False, + default="bounce", + choices=["orbit", "bounce"], + ) + parser.add_argument("--detector-speed", type=float, required=False, default=10.0) + parser.add_argument("--source-speed", type=float, required=False, default=0.0) + parser.add_argument("--sigma", type=float, required=False, default=1.0) + parser.add_argument("--time-steps", type=int, required=False, default=100) + parser.add_argument("--time-interval", type=float, required=False, default=0.3) + parser.add_argument("--samples-per-snapshot", type=int, required=False, default=3) + parser.add_argument("--sessions", type=int, required=False, default=1024) + parser.add_argument( + "--output", type=str, required=False, default="sessions-default" + ) + parser.add_argument("--reference", type=bool, required=False, default=False) + parser.add_argument("--cpus", type=int, required=False, default=8) + parser.add_argument("--live", type=bool, required=False, default=False) + parser.add_argument("--profile", type=bool, required=False, default=False) + parser.add_argument("--fixed-detector", type=float, required=False, nargs="+") - args = parser.parse_args() + args = parser.parse_args() - os.mkdir(args.output) - dump(args,"/".join([args.output,'args.pkl']),compression="lzma") - if not args.live: - if args.profile: - #import cProfile, pstats, io - #from pstats import SortKey - #pr = cProfile.Profile() - #pr.enable() - for session_idx in np.arange(args.sessions): - result = generate_session((args,session_idx)) - #pr.disable() - #s = io.StringIO() - #sortby = SortKey.CUMULATIVE - #ps = pstats.Stats(pr, stream=s).sort_stats(sortby) - #ps.print_stats() - #print(s.getvalue()) - else: - result = Parallel(n_jobs=args.cpus)(delayed(generate_session_and_dump)((args,session_idx)) for session_idx in tqdm(range(args.sessions))) + os.mkdir(args.output) + dump(args, "/".join([args.output, "args.pkl"]), compression="lzma") + if not args.live: + if args.profile: + # import cProfile, pstats, io + # from pstats import SortKey + # pr = cProfile.Profile() + # pr.enable() + for session_idx in np.arange(args.sessions): + result = generate_session((args, session_idx)) + # pr.disable() + # s = io.StringIO() + # sortby = SortKey.CUMULATIVE + # ps = pstats.Stats(pr, stream=s).sort_stats(sortby) + # ps.print_stats() + # print(s.getvalue()) + else: + result = Parallel(n_jobs=args.cpus)( + delayed(generate_session_and_dump)((args, session_idx)) + for session_idx in tqdm(range(args.sessions)) + ) diff --git a/software/model_training_and_inference/02_check_data.py b/software/model_training_and_inference/02_check_data.py index 04b08503..9da048fc 100644 --- a/software/model_training_and_inference/02_check_data.py +++ b/software/model_training_and_inference/02_check_data.py @@ -1,4 +1,3 @@ - import os import pickle from functools import cache @@ -8,95 +7,106 @@ import torch from torch.utils.data import DataLoader, Dataset -from utils.image_utils import (detector_positions_to_theta_grid, - labels_to_source_images, radio_to_image) +from utils.image_utils import ( + detector_positions_to_theta_grid, + labels_to_source_images, + radio_to_image, +) from utils.plot import plot_space from utils.spf_dataset import * -def plot_single_session(d): - fig=plt.figure(figsize=(16,4)) - axs=fig.subplots(1,4) - det_x,det_y=d['detector_position_at_t'][0] #*self.width - detector_pos_str="Det (X=%d,Y=%d)" % (det_x,det_y) - - axs[0].imshow(d['source_image_at_t'][0,0].T) - axs[0].set_title("Source emitters") - - axs[1].imshow(d['detector_theta_image_at_t'][0,0].T) - axs[1].set_title("Detector theta angle") - - axs[2].imshow(d['radio_image_at_t'][0,0].T) - axs[2].set_title("Radio feature (steering)") - - for idx in range(3): - axs[idx].set_xlabel("X") - axs[idx].set_ylabel("Y") - - axs[3].plot(d['thetas_at_t'][0],d['beam_former_outputs_at_t'][0]) - axs[3].set_title("Radio feature (steering)") - axs[3].set_xlabel("Theta") - axs[3].set_ylabel("Power") - fig.suptitle(detector_pos_str) - fig.tight_layout() - fig.show() - breakpoint() - -if __name__=='__main__': - sessions_dir='./sessions-default' - #test task1 - if False: - #load a dataset - ds=SessionsDatasetTask1(sessions_dir) - ds[253] - ds=SessionsDataset(sessions_dir) - - r_idxs=np.arange(len(ds)) - np.random.shuffle(r_idxs) - fig=plt.figure(figsize=(8,8)) - for r_idx in r_idxs[:4]: - fig.clf() - axs=fig.subplots(3,3) - x=ds[r_idx] - theta_at_pos=detector_positions_to_theta_grid(x['detector_position_at_t'][None],ds.args.width) # batch, snapshot, 1, x ,y - dist_at_pos=detector_positions_to_distance(x['detector_position_at_t'][None],ds.args.width) # batch, snapshot, 1, x ,y - theta_idxs=(((theta_at_pos+np.pi)/(2*np.pi))*257) #.int().float() #~ [2, 8, 1, 128, 128]) - - label_images=labels_to_source_images(x['source_positions_at_t'][None],ds.args.width) - - _,s,_,_x,_y=dist_at_pos.shape - for s_idx in np.arange(s): - m=theta_idxs[0,s_idx] - for idx in np.arange(257): - m[m==idx]=np.abs(x['beam_former_outputs_at_t'][s_idx,idx]) #.abs() #.log() - for idx in np.arange(3): - axs[idx,0].imshow(theta_idxs[0,idx,0]) - axs[idx,1].imshow(label_images[0,idx,0]) - axs[idx,2].imshow( (theta_idxs[0,:idx+1,0].mean(axis=0)) ) - plt.pause(1) - #plot the space diagram for some samples - fig,ax=plt.subplots(2,2,figsize=(8,8)) - [ plot_space(ax[i//2,i%2], ds[r_idxs[i]]) for i in np.arange(4) ] - plt.title("Task1") - plt.show() - - #test task2 - if True: - #load a dataset - ds=SessionsDatasetTask2(sessions_dir) - ds[253] - plot_single_session(ds[0]) - plot_single_session(ds[200]) - plot_single_session(ds[100]) - ds=SessionsDataset(sessions_dir) - - #plot the space diagram for some samples - fig,ax=plt.subplots(2,2,figsize=(8,8)) - r_idxs=np.arange(len(ds)) - np.random.shuffle(r_idxs) - [ plot_space(ax[i//2,i%2], ds[r_idxs[i]]) for i in np.arange(4) ] - plt.title("Task2") - plt.show() - - - +def plot_single_session(d): + fig = plt.figure(figsize=(16, 4)) + axs = fig.subplots(1, 4) + det_x, det_y = d["detector_position_at_t"][0] # *self.width + detector_pos_str = "Det (X=%d,Y=%d)" % (det_x, det_y) + + axs[0].imshow(d["source_image_at_t"][0, 0].T) + axs[0].set_title("Source emitters") + + axs[1].imshow(d["detector_theta_image_at_t"][0, 0].T) + axs[1].set_title("Detector theta angle") + + axs[2].imshow(d["radio_image_at_t"][0, 0].T) + axs[2].set_title("Radio feature (steering)") + + for idx in range(3): + axs[idx].set_xlabel("X") + axs[idx].set_ylabel("Y") + + axs[3].plot(d["thetas_at_t"][0], d["beam_former_outputs_at_t"][0]) + axs[3].set_title("Radio feature (steering)") + axs[3].set_xlabel("Theta") + axs[3].set_ylabel("Power") + fig.suptitle(detector_pos_str) + fig.tight_layout() + fig.show() + breakpoint() + + +if __name__ == "__main__": + sessions_dir = "./sessions-default" + # test task1 + if False: + # load a dataset + ds = SessionsDatasetTask1(sessions_dir) + ds[253] + ds = SessionsDataset(sessions_dir) + + r_idxs = np.arange(len(ds)) + np.random.shuffle(r_idxs) + fig = plt.figure(figsize=(8, 8)) + for r_idx in r_idxs[:4]: + fig.clf() + axs = fig.subplots(3, 3) + x = ds[r_idx] + theta_at_pos = detector_positions_to_theta_grid( + x["detector_position_at_t"][None], ds.args.width + ) # batch, snapshot, 1, x ,y + dist_at_pos = detector_positions_to_distance( + x["detector_position_at_t"][None], ds.args.width + ) # batch, snapshot, 1, x ,y + theta_idxs = ( + (theta_at_pos + np.pi) / (2 * np.pi) + ) * 257 # .int().float() #~ [2, 8, 1, 128, 128]) + + label_images = labels_to_source_images( + x["source_positions_at_t"][None], ds.args.width + ) + + _, s, _, _x, _y = dist_at_pos.shape + for s_idx in np.arange(s): + m = theta_idxs[0, s_idx] + for idx in np.arange(257): + m[m == idx] = np.abs( + x["beam_former_outputs_at_t"][s_idx, idx] + ) # .abs() #.log() + for idx in np.arange(3): + axs[idx, 0].imshow(theta_idxs[0, idx, 0]) + axs[idx, 1].imshow(label_images[0, idx, 0]) + axs[idx, 2].imshow((theta_idxs[0, : idx + 1, 0].mean(axis=0))) + plt.pause(1) + # plot the space diagram for some samples + fig, ax = plt.subplots(2, 2, figsize=(8, 8)) + [plot_space(ax[i // 2, i % 2], ds[r_idxs[i]]) for i in np.arange(4)] + plt.title("Task1") + plt.show() + + # test task2 + if True: + # load a dataset + ds = SessionsDatasetTask2(sessions_dir) + ds[253] + plot_single_session(ds[0]) + plot_single_session(ds[200]) + plot_single_session(ds[100]) + ds = SessionsDataset(sessions_dir) + + # plot the space diagram for some samples + fig, ax = plt.subplots(2, 2, figsize=(8, 8)) + r_idxs = np.arange(len(ds)) + np.random.shuffle(r_idxs) + [plot_space(ax[i // 2, i % 2], ds[r_idxs[i]]) for i in np.arange(4)] + plt.title("Task2") + plt.show() diff --git a/software/model_training_and_inference/12_task2_model_training.py b/software/model_training_and_inference/12_task2_model_training.py index 1dc65e07..4e83ec98 100644 --- a/software/model_training_and_inference/12_task2_model_training.py +++ b/software/model_training_and_inference/12_task2_model_training.py @@ -14,521 +14,746 @@ from torch.utils.data import dataset, random_split from utils.image_utils import labels_to_source_images -from models.models import (SingleSnapshotNet, SnapshotNet, Task1Net, TransformerEncOnlyModel, - UNet) -from utils.spf_dataset import SessionsDataset, SessionsDatasetTask2, collate_fn, output_cols, input_cols - -torch.set_printoptions(precision=5,sci_mode=False,linewidth=1000) - - -def src_pos_from_radial(inputs,outputs): - det_pos=inputs[:,:,input_cols['det_pos']] - - theta=outputs[:,:,[cols_for_loss.index(output_cols['src_theta'][0])]]*np.pi - dist=outputs[:,:,[cols_for_loss.index(output_cols['src_dist'][0])]] - - theta=theta.float() - dist=dist.float() - det_pos=det_pos.float() - return torch.stack([torch.sin(theta),torch.cos(theta)],axis=2)[...,0]*dist+det_pos - -def model_forward(d_model,data,args,train_test_label,update,plot=True): - #if radio_inputs.isnan().any(): - # breakpoint() - d_model['optimizer'].zero_grad() - use_all_data=True - if use_all_data: - b,s,_=data['inputs']['radio_feature'].shape - model_s=d_model['snapshots_per_sample'] - #assert(s%model_s==0) - _data={ - 'inputs':{ - 'drone_state':data['inputs']['drone_state'].reshape(b*s//model_s,model_s,-1), - 'radio_feature':data['inputs']['radio_feature'].reshape(b*s//model_s,model_s,-1) - }, - 'labels':data['labels'].reshape(b*s//model_s,model_s,-1) - } - else: - _data={ - 'inputs':{ - 'drone_state':data['inputs']['drone_state'][:,:d_model['snapshots_per_sample']], - 'radio_feature':data['inputs']['radio_feature'][:,:d_model['snapshots_per_sample']] - }, - 'labels':data['labels'][:,:d_model['snapshots_per_sample']] - } - #_radio_images=radio_images[:,:d_model['snapshots_per_sample'],0] # reshape to b,s,w,h - #_label_images=label_images[:,d_model['snapshots_per_sample']-1] # reshape to b,1,w,h - losses={} - - preds=d_model['model'](_data['inputs'])#.detach().cpu() - assert(not _data['inputs']['radio_feature'].isnan().any()) - assert(not _data['inputs']['drone_state'].isnan().any()) - #for k in preds: - # preds[k]=preds[k].detach().cpu() - transformer_loss=0.0 - single_snapshot_loss=0.0 - fc_loss=0.0 - if 'transformer_pred' in preds: - if preds['transformer_pred'].mean(axis=1).var(axis=0).mean().item()<1e-13: - d_model['dead']=True - else: - d_model['dead']=False - transformer_loss = criterion(preds['transformer_pred'],_data['labels']) - losses['transformer_loss']=transformer_loss.detach().item() - losses['transformer_stats']=(preds['transformer_pred']-_data['labels']).pow(2).mean(axis=[0,1]).detach().cpu() - _p=preds['transformer_pred'].detach().cpu() - assert(not preds['transformer_pred'].isnan().any()) - if 'single_snapshot_pred' in preds: - single_snapshot_loss = criterion(preds['single_snapshot_pred'],_data['labels']) - losses['single_snapshot_loss']=single_snapshot_loss.item() - losses['single_snapshot_stats']=(preds['single_snapshot_pred']-_data['labels']).pow(2).mean(axis=[0,1]).detach().cpu() - elif 'fc_pred' in preds: - fc_loss = criterion(preds['fc_pred'],_data['labels'][:,[-1]]) - losses['fc_loss']=fc_loss.item() - _p=preds['fc_pred'].detach().cpu() - #losses['fc_stats']=(preds['fc_pred']-_labels).pow(2).mean(axis=[0,1]).cpu() - if plot and (update%args.plot_every)==args.plot_every-1: - d_model['fig'].clf() - _ri=_data['inputs']['drone_state'].detach().cpu() - _l=_data['labels'].detach().cpu() - axs=d_model['fig'].subplots(1,4,sharex=True,sharey=True) - d_model['fig'].suptitle(d_model['name']) - axs[0].scatter(_ri[0,:,input_cols['det_pos'][0]],_ri[0,:,input_cols['det_pos'][1]],label='detector positions',s=1) - if 'src_pos' in args.losses: - src_pos_idxs=[] - for idx in range(2): - src_pos_idxs.append(cols_for_loss.index(output_cols['src_pos'][idx])) - axs[1].scatter(_l[0,:,src_pos_idxs[0]],_l[0,:,src_pos_idxs[1]],label='source positions',c='r') - - axs[2].scatter(_l[0,:,src_pos_idxs[0]],_l[0,:,src_pos_idxs[1]],label='real positions',c='b',alpha=0.1,s=7) - axs[2].scatter(_p[0,:,src_pos_idxs[0]],_p[0,:,src_pos_idxs[1]],label='predicted positions',c='r',alpha=0.3,s=7) - axs[2].legend() - if 'src_theta' in args.losses and 'src_dist' in args.losses: - pos_from_preds=src_pos_from_radial(_ri,_l) - axs[3].scatter(pos_from_preds[0,:,0],pos_from_preds[0,:,1],label='real radial positions',c='b',alpha=0.1,s=7) - pos_from_preds=src_pos_from_radial(_ri,_p) - axs[3].scatter(pos_from_preds[0,:,0],pos_from_preds[0,:,1],label='predicted radial positions',c='r',alpha=0.3,s=7) - axs[3].legend() - - for idx in [0,1,2]: - axs[idx].legend() - axs[idx].set_xlim([-1,1]) - axs[idx].set_ylim([-1,1]) - d_model['fig'].savefig('%s%s_%d_%s.png' % (args.output_prefix,d_model['name'],update,train_test_label)) - d_model['fig'].canvas.draw_idle() - loss=(1.0-args.transformer_loss_balance)*transformer_loss+args.transformer_loss_balance*single_snapshot_loss+fc_loss - if ikeep_n_saves: - fn=saves.pop(0) - if os.path.exists(fn): - os.remove(fn) - -def plot_loss(running_losses, - baseline_loss, - baseline_image_loss, - xtick_spacing, - mean_chunk, - output_prefix, - fig, - title, - update): - fig.clf() - fig.suptitle(title) - axs=fig.subplots(1,4,sharex=True) - axs[1].sharex(axs[0]) - axs[2].sharex(axs[0]) - xs=np.arange(len(baseline_loss['baseline']))*xtick_spacing - for i in range(3): - axs[i].plot(xs,baseline_loss['baseline'],label='baseline') - axs[i].set_xlabel("time") - axs[i].set_ylabel("log loss") - if baseline_image_loss is not None: - axs[3].plot(xs,baseline_image_loss['baseline_image'],label='baseline image') - axs[0].set_title("Transformer loss") - axs[1].set_title("Single snapshot loss") - axs[2].set_title("FC loss") - axs[3].set_title("Image loss") - for d_model in models: - losses=model_to_losses(running_losses[d_model['name']],mean_chunk) - if 'transformer_loss' in losses: - axs[0].plot(xs,losses['transformer_loss'],label=d_model['name']) - if 'single_snapshot_loss' in losses: - axs[1].plot(xs,losses['single_snapshot_loss'],label=d_model['name']) - if 'fc_loss' in losses: - axs[2].plot(xs,losses['fc_loss'],label=d_model['name']) - if 'image_loss' in losses: - axs[3].plot(xs,losses['image_loss'],label=d_model['name']) - for i in range(4): - axs[i].legend() - fig.tight_layout() - fig.savefig('%sloss_%s_%d.png' % (output_prefix,title,update)) - fig.canvas.draw_idle() - -if __name__=='__main__': - parser = argparse.ArgumentParser() - parser.add_argument('--device', type=str, required=False, default='cpu') - parser.add_argument('--embedding-warmup', type=int, required=False, default=0) - parser.add_argument('--snapshots-per-sample', type=int, required=False, default=[1,4,8], nargs="+") - parser.add_argument('--print-every', type=int, required=False, default=100) - parser.add_argument('--lr-scheduler-every', type=int, required=False, default=256) - parser.add_argument('--plot-every', type=int, required=False, default=1024) - parser.add_argument('--save-every', type=int, required=False, default=1000) - parser.add_argument('--test-mbs', type=int, required=False, default=8) - parser.add_argument('--output-prefix', type=str, required=False, default='model_out') - parser.add_argument('--test-fraction', type=float, required=False, default=0.2) - parser.add_argument('--weight-decay', type=float, required=False, default=0.0) - parser.add_argument('--transformer-loss-balance', type=float, required=False, default=0.1) - parser.add_argument('--type', type=str, required=False, default=32) - parser.add_argument('--seed', type=int, required=False, default=0) - parser.add_argument('--keep-n-saves', type=int, required=False, default=2) - parser.add_argument('--epochs', type=int, required=False, default=20000) - parser.add_argument('--positional-encoding-len', type=int, required=False, default=0) - parser.add_argument('--mb', type=int, required=False, default=64) - parser.add_argument('--workers', type=int, required=False, default=4) - parser.add_argument('--dataset', type=str, required=False, default='./sessions-default') - parser.add_argument('--lr-image', type=float, required=False, default=0.05) - parser.add_argument('--lr-direct', type=float, required=False, default=0.01) - parser.add_argument('--lr-transformer', type=float, required=False, default=0.00001) - parser.add_argument('--plot', type=bool, required=False, default=False) - parser.add_argument('--transformer-input', type=str, required=False, default=['drone_state','embedding','single_snapshot_pred'],nargs="+") - parser.add_argument('--transformer-dmodel', type=int, required=False, default=128+256) - parser.add_argument('--clip', type=float, required=False, default=0.5) - parser.add_argument('--losses', type=str, required=False, default="src_pos,src_theta,src_dist") #,src_theta,src_dist,det_delta,det_theta,det_space") - args = parser.parse_args() - - dtype=torch.float32 - if args.type=='16': - dtype=torch.float16 - - if args.plot==False: - import matplotlib - matplotlib.use('Agg') - import matplotlib.pyplot as plt - - - torch.manual_seed(args.seed) - random.seed(args.seed) - np.random.seed(args.seed) - - start_time=time.time() - - #lets see if output dir exists , if not make it - basename=os.path.basename(args.output_prefix) - if not os.path.exists(basename): - os.makedirs(basename) - - - cols_for_loss=[] - for k in output_cols: - if k in args.losses.split(','): - cols_for_loss+=output_cols[k] - - device=torch.device(args.device) - print("init dataset") - ds=SessionsDatasetTask2(args.dataset,snapshots_in_sample=max(args.snapshots_per_sample)) - #ds_test=SessionsDatasetTask2(args.test_dataset,snapshots_in_sample=max(args.snapshots_per_sample)) - train_size=int(len(ds)*args.test_fraction) - test_size=len(ds)-train_size - - #need to generate separate files for this not to train test leak - #ds_train, ds_test = random_split(ds, [1-args.test_fraction, args.test_fraction]) - - ds_train = torch.utils.data.Subset(ds, np.arange(train_size)) - ds_test = torch.utils.data.Subset(ds, np.arange(train_size, train_size + test_size)) - - print("init dataloader") - trainloader = torch.utils.data.DataLoader( - ds_train, - batch_size=args.mb, - shuffle=True, - num_workers=args.workers, - collate_fn=collate_fn) - testloader = torch.utils.data.DataLoader( - ds_test, - batch_size=args.mb, - shuffle=True, - num_workers=args.workers, - collate_fn=collate_fn) - - print("init network") - models=[] - if True: - for n_layers in [2,4,8,16,32]:#,32,64]: #,32,64]: - for snapshots_per_sample in args.snapshots_per_sample: - if snapshots_per_sample==1: - continue - models.append( - { - 'name':'%d snapshots (l%d)' % (snapshots_per_sample,n_layers), - 'model':SnapshotNet( - snapshots_per_sample, - n_layers=n_layers, - d_model=args.transformer_dmodel, - n_outputs=len(cols_for_loss), - ssn_n_outputs=len(cols_for_loss), - dropout=0.0, - positional_encoding_len=args.positional_encoding_len, - tformer_input=args.transformer_input), - 'snapshots_per_sample':snapshots_per_sample, - 'images':False, - 'lr':args.lr_transformer, - 'images':False,'fig':plt.figure(figsize=(18,4)), - 'dead':False, - } - ) - if False: - for snapshots_per_sample in args.snapshots_per_sample: - models.append( - { - 'name':'task1net%d' % snapshots_per_sample, - 'model':Task1Net(266*snapshots_per_sample,n_outputs=len(cols_for_loss)), - 'snapshots_per_sample':snapshots_per_sample, - 'images':False,'fig':plt.figure(figsize=(18,4)), - 'lr':args.lr_direct, - 'dead':False, - } - ) - if True: - for snapshots_per_sample in args.snapshots_per_sample: - models.append( - { - 'name':'FCNet %d' % snapshots_per_sample, - 'model':SingleSnapshotNet(d_input_feature=266, - d_hid=64, - d_embed=64, - n_layers=4, - n_outputs=len(cols_for_loss), - dropout=0.0, - snapshots_per_sample=snapshots_per_sample), - 'snapshots_per_sample':snapshots_per_sample, - 'images':False,'fig':plt.figure(figsize=(18,4)), - 'lr':args.lr_direct, - 'dead':False - } - ) - - if False: - for snapshots_per_sample in args.snapshots_per_sample: - models.append( - { - 'name':'Unet %d' % snapshots_per_sample, - 'model':UNet(in_channels=snapshots_per_sample,out_channels=1,width=128), - 'snapshots_per_sample':snapshots_per_sample, - 'images':True,'fig':plt.figure(figsize=(14,4)), - 'normalize_input':True, - 'lr':args.lr_image, - 'dead':False - } - ) - - using_images=False - for d_model in models: - if d_model['images']: - using_images=True - - - #move the models to the device - for d_net in models: - d_net['model']=d_net['model'].to(dtype).to(device) - loss_figs={ - 'train':plt.figure(figsize=(14*3,6)), - 'test':plt.figure(figsize=(14*3,6))} - - for d_model in models: - d_model['optimizer']=optim.Adam(d_model['model'].parameters(),lr=d_model['lr'],weight_decay=args.weight_decay) - d_model['scheduler']=optim.lr_scheduler.LinearLR( - d_model['optimizer'], - start_factor=0.001, - end_factor=1.0, - total_iters=30, - verbose=False) - - criterion = nn.MSELoss().to(device) - - print("training loop") - running_losses={'train':{},'test':{}} - for k in ['train','test']: - running_losses[k]={ d['name']:[] for d in models} - running_losses[k]['baseline']=[] - running_losses[k]['baseline_image']=[] - - saves=[] - - - def prep_data(data): - prepared_data={ - 'inputs':{ - 'radio_feature':data['inputs']['radio_feature'].to(dtype).to(device), - 'drone_state':data['inputs']['drone_state'].to(dtype).to(device) - }, - 'labels':data['labels'][...,cols_for_loss].to(dtype).to(device) - } - - assert(not data['inputs']['radio_feature'].isnan().any()) - assert(not data['inputs']['radio_feature'].isinf().any()) - assert(not data['inputs']['drone_state'].isnan().any()) - assert(not data['inputs']['drone_state'].isinf().any()) - - return prepared_data - - test_iterator = iter(testloader) - for epoch in range(args.epochs): - for i, data in enumerate(trainloader, 0): - #move to device, do final prep - prepared_data=prep_data(data) - if True: #torch.cuda.amp.autocast(): - for d_model in models: - for p in d_net['model'].parameters(): - if p.isnan().any(): - breakpoint() - loss,losses=model_forward( - d_model, - prepared_data, - args, - 'train', - update=i, - plot=True) - if i%args.lr_scheduler_every==args.lr_scheduler_every-1: - d_model['scheduler'].step() - loss.backward() - running_losses['train'][d_model['name']].append(losses) - if args.clip>0: - torch.nn.utils.clip_grad_norm_(d_net['model'].parameters(), args.clip) # clip gradients - d_model['optimizer'].step() - for p in d_net['model'].parameters(): - if p.isnan().any(): - breakpoint() - labels=prepared_data['labels'] - running_losses['train']['baseline'].append( {'baseline':criterion(labels*0+labels.mean(axis=[0,1],keepdim=True), labels).item() } ) - if using_images: - running_losses['train']['baseline_image'].append( {'baseline_image':criterion(label_images*0+label_images.mean(), label_images).item() } ) - - if i%args.print_every==args.print_every-1: - for idx in np.arange(args.test_mbs): - try: - data = next(test_iterator) - except StopIteration: - test_iterator = iter(testloader) - data = next(test_iterator) - prepared_data=prep_data(data) - with torch.no_grad(): - for d_model in models: - loss,losses=model_forward( - d_model, - prepared_data, - args, - 'test', - update=i, - plot=idx==0) - running_losses['test'][d_model['name']].append(losses) - labels=prepared_data['labels'] - running_losses['test']['baseline'].append( {'baseline':criterion(labels*0+labels.mean(axis=[0,1],keepdim=True), labels).item() } ) - if using_images: - running_losses['test']['baseline_image'].append( {'baseline_image':criterion(label_images*0+label_images.mean(), label_images).item() } ) - - - if i==0 or i%args.save_every==args.save_every-1: - save(args,running_losses,models,i,args.keep_n_saves) - - if i % args.print_every == args.print_every-1: - - train_baseline_loss=model_to_losses(running_losses['train']['baseline'],args.print_every) - test_baseline_loss=model_to_losses(running_losses['test']['baseline'],args.test_mbs) - train_baseline_image_loss=None - test_baseline_image_loss=None - print(f'[{epoch + 1}, {i + 1:5d}]') - if using_images: - train_baseline_image_loss=model_to_losses(running_losses['train']['baseline_image'],args.print_every) - test_baseline_image_loss=model_to_losses(running_losses['test']['baseline_image'],args.test_mbs) - print(f'\tTrain: baseline: {train_baseline_loss["baseline"][-1]:.3f}, baseline_image: {train_baseline_image_loss["baseline_image"][-1]:.3f} , time { (time.time()-start_time)/(i+1) :.3f} / batch' ) - print(f'\tTest: baseline: {test_baseline_loss["baseline"][-1]:.3f}, baseline_image: {test_baseline_image_loss["baseline_image"][-1]:.3f} , time { (time.time()-start_time)/(i+1) :.3f} / batch' ) - else: - print(f'\tTrain: baseline: {train_baseline_loss["baseline"][-1]:.3f} , time { (time.time()-start_time)/(i+1) :.3f} / batch' ) - print(f'\tTest: baseline: {test_baseline_loss["baseline"][-1]:.3f}, time { (time.time()-start_time)/(i+1) :.3f} / batch' ) - loss_str="\t"+"\n\t".join( - [ "%s(%s):(tr)%s,(ts)%s" % (d['name'],str(d['dead']), - model_to_loss_str(running_losses['train'][d['name']],args.print_every), - model_to_loss_str(running_losses['test'][d['name']],args.test_mbs) - ) for d in models ]) - print(loss_str) - loss_str="\n".join( - [ "\t%s:(tr)\n%s\n\t%s:(ts)\n%s" % (d['name'], - model_to_stats_str(running_losses['train'][d['name']],args.print_every), - d['name'], - model_to_stats_str(running_losses['test'][d['name']],args.test_mbs) - ) for d in models ]) - print("\t\t\t%s" % stats_title()) - print(loss_str) - if i//args.print_every>2 and i % args.plot_every == args.plot_every-1: - plot_loss(running_losses=running_losses['train'], - baseline_loss=train_baseline_loss, - baseline_image_loss=train_baseline_image_loss, - xtick_spacing=args.print_every, - mean_chunk=args.print_every, - output_prefix=args.output_prefix, - fig=loss_figs['train'], - title='Train',update=i) - plot_loss(running_losses=running_losses['test'], - baseline_loss=test_baseline_loss, - baseline_image_loss=test_baseline_image_loss, - xtick_spacing=args.print_every, - mean_chunk=args.test_mbs, - output_prefix=args.output_prefix, - fig=loss_figs['test'], - title='Test',update=i) - if args.plot: - plt.pause(0.5) - - print('Finished Training') # but do we ever really get here? + title_str = [] + for col in args.losses.split(","): + for _ in range(len(output_cols[col])): + title_str.append(col) + return "\t".join(title_str) + + +def model_to_stats_str(running_loss, mean_chunk): + if len(running_loss) == 0: + return "" + losses = model_to_losses(running_loss, mean_chunk) + loss_str = [] + for k in ["transformer_stats", "single_snapshot_stats"]: + if k in losses: + loss_str.append( + "\t\t%s\t%s" + % (k, "\t".join(["%0.4f" % v.item() for v in losses[k][-1]])) + ) + return "\n".join(loss_str) + + +def save(args, running_losses, models, iteration, keep_n_saves): + fn = "%ssave_%d.pkl" % (args.output_prefix, iteration) + pickle.dump( + {"models": models, "args": args, "running_losses": running_losses}, + open(fn, "wb"), + ) + saves.append(fn) + while len(saves) > keep_n_saves: + fn = saves.pop(0) + if os.path.exists(fn): + os.remove(fn) + + +def plot_loss( + running_losses, + baseline_loss, + baseline_image_loss, + xtick_spacing, + mean_chunk, + output_prefix, + fig, + title, + update, +): + fig.clf() + fig.suptitle(title) + axs = fig.subplots(1, 4, sharex=True) + axs[1].sharex(axs[0]) + axs[2].sharex(axs[0]) + xs = np.arange(len(baseline_loss["baseline"])) * xtick_spacing + for i in range(3): + axs[i].plot(xs, baseline_loss["baseline"], label="baseline") + axs[i].set_xlabel("time") + axs[i].set_ylabel("log loss") + if baseline_image_loss is not None: + axs[3].plot(xs, baseline_image_loss["baseline_image"], label="baseline image") + axs[0].set_title("Transformer loss") + axs[1].set_title("Single snapshot loss") + axs[2].set_title("FC loss") + axs[3].set_title("Image loss") + for d_model in models: + losses = model_to_losses(running_losses[d_model["name"]], mean_chunk) + if "transformer_loss" in losses: + axs[0].plot(xs, losses["transformer_loss"], label=d_model["name"]) + if "single_snapshot_loss" in losses: + axs[1].plot(xs, losses["single_snapshot_loss"], label=d_model["name"]) + if "fc_loss" in losses: + axs[2].plot(xs, losses["fc_loss"], label=d_model["name"]) + if "image_loss" in losses: + axs[3].plot(xs, losses["image_loss"], label=d_model["name"]) + for i in range(4): + axs[i].legend() + fig.tight_layout() + fig.savefig("%sloss_%s_%d.png" % (output_prefix, title, update)) + fig.canvas.draw_idle() + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--device", type=str, required=False, default="cpu") + parser.add_argument("--embedding-warmup", type=int, required=False, default=0) + parser.add_argument( + "--snapshots-per-sample", type=int, required=False, default=[1, 4, 8], nargs="+" + ) + parser.add_argument("--print-every", type=int, required=False, default=100) + parser.add_argument("--lr-scheduler-every", type=int, required=False, default=256) + parser.add_argument("--plot-every", type=int, required=False, default=1024) + parser.add_argument("--save-every", type=int, required=False, default=1000) + parser.add_argument("--test-mbs", type=int, required=False, default=8) + parser.add_argument( + "--output-prefix", type=str, required=False, default="model_out" + ) + parser.add_argument("--test-fraction", type=float, required=False, default=0.2) + parser.add_argument("--weight-decay", type=float, required=False, default=0.0) + parser.add_argument( + "--transformer-loss-balance", type=float, required=False, default=0.1 + ) + parser.add_argument("--type", type=str, required=False, default=32) + parser.add_argument("--seed", type=int, required=False, default=0) + parser.add_argument("--keep-n-saves", type=int, required=False, default=2) + parser.add_argument("--epochs", type=int, required=False, default=20000) + parser.add_argument( + "--positional-encoding-len", type=int, required=False, default=0 + ) + parser.add_argument("--mb", type=int, required=False, default=64) + parser.add_argument("--workers", type=int, required=False, default=4) + parser.add_argument( + "--dataset", type=str, required=False, default="./sessions-default" + ) + parser.add_argument("--lr-image", type=float, required=False, default=0.05) + parser.add_argument("--lr-direct", type=float, required=False, default=0.01) + parser.add_argument("--lr-transformer", type=float, required=False, default=0.00001) + parser.add_argument("--plot", type=bool, required=False, default=False) + parser.add_argument( + "--transformer-input", + type=str, + required=False, + default=["drone_state", "embedding", "single_snapshot_pred"], + nargs="+", + ) + parser.add_argument( + "--transformer-dmodel", type=int, required=False, default=128 + 256 + ) + parser.add_argument("--clip", type=float, required=False, default=0.5) + parser.add_argument( + "--losses", type=str, required=False, default="src_pos,src_theta,src_dist" + ) # ,src_theta,src_dist,det_delta,det_theta,det_space") + args = parser.parse_args() + + dtype = torch.float32 + if args.type == "16": + dtype = torch.float16 + + if args.plot == False: + import matplotlib + + matplotlib.use("Agg") + import matplotlib.pyplot as plt + + torch.manual_seed(args.seed) + random.seed(args.seed) + np.random.seed(args.seed) + + start_time = time.time() + + # lets see if output dir exists , if not make it + basename = os.path.basename(args.output_prefix) + if not os.path.exists(basename): + os.makedirs(basename) + + cols_for_loss = [] + for k in output_cols: + if k in args.losses.split(","): + cols_for_loss += output_cols[k] + + device = torch.device(args.device) + print("init dataset") + ds = SessionsDatasetTask2( + args.dataset, snapshots_in_sample=max(args.snapshots_per_sample) + ) + # ds_test=SessionsDatasetTask2(args.test_dataset,snapshots_in_sample=max(args.snapshots_per_sample)) + train_size = int(len(ds) * args.test_fraction) + test_size = len(ds) - train_size + + # need to generate separate files for this not to train test leak + # ds_train, ds_test = random_split(ds, [1-args.test_fraction, args.test_fraction]) + + ds_train = torch.utils.data.Subset(ds, np.arange(train_size)) + ds_test = torch.utils.data.Subset(ds, np.arange(train_size, train_size + test_size)) + + print("init dataloader") + trainloader = torch.utils.data.DataLoader( + ds_train, + batch_size=args.mb, + shuffle=True, + num_workers=args.workers, + collate_fn=collate_fn, + ) + testloader = torch.utils.data.DataLoader( + ds_test, + batch_size=args.mb, + shuffle=True, + num_workers=args.workers, + collate_fn=collate_fn, + ) + + print("init network") + models = [] + if True: + for n_layers in [2, 4, 8, 16, 32]: # ,32,64]: #,32,64]: + for snapshots_per_sample in args.snapshots_per_sample: + if snapshots_per_sample == 1: + continue + models.append( + { + "name": "%d snapshots (l%d)" % (snapshots_per_sample, n_layers), + "model": SnapshotNet( + snapshots_per_sample, + n_layers=n_layers, + d_model=args.transformer_dmodel, + n_outputs=len(cols_for_loss), + ssn_n_outputs=len(cols_for_loss), + dropout=0.0, + positional_encoding_len=args.positional_encoding_len, + tformer_input=args.transformer_input, + ), + "snapshots_per_sample": snapshots_per_sample, + "images": False, + "lr": args.lr_transformer, + "images": False, + "fig": plt.figure(figsize=(18, 4)), + "dead": False, + } + ) + if False: + for snapshots_per_sample in args.snapshots_per_sample: + models.append( + { + "name": "task1net%d" % snapshots_per_sample, + "model": Task1Net( + 266 * snapshots_per_sample, n_outputs=len(cols_for_loss) + ), + "snapshots_per_sample": snapshots_per_sample, + "images": False, + "fig": plt.figure(figsize=(18, 4)), + "lr": args.lr_direct, + "dead": False, + } + ) + if True: + for snapshots_per_sample in args.snapshots_per_sample: + models.append( + { + "name": "FCNet %d" % snapshots_per_sample, + "model": SingleSnapshotNet( + d_input_feature=266, + d_hid=64, + d_embed=64, + n_layers=4, + n_outputs=len(cols_for_loss), + dropout=0.0, + snapshots_per_sample=snapshots_per_sample, + ), + "snapshots_per_sample": snapshots_per_sample, + "images": False, + "fig": plt.figure(figsize=(18, 4)), + "lr": args.lr_direct, + "dead": False, + } + ) + + if False: + for snapshots_per_sample in args.snapshots_per_sample: + models.append( + { + "name": "Unet %d" % snapshots_per_sample, + "model": UNet( + in_channels=snapshots_per_sample, out_channels=1, width=128 + ), + "snapshots_per_sample": snapshots_per_sample, + "images": True, + "fig": plt.figure(figsize=(14, 4)), + "normalize_input": True, + "lr": args.lr_image, + "dead": False, + } + ) + + using_images = False + for d_model in models: + if d_model["images"]: + using_images = True + + # move the models to the device + for d_net in models: + d_net["model"] = d_net["model"].to(dtype).to(device) + loss_figs = { + "train": plt.figure(figsize=(14 * 3, 6)), + "test": plt.figure(figsize=(14 * 3, 6)), + } + + for d_model in models: + d_model["optimizer"] = optim.Adam( + d_model["model"].parameters(), + lr=d_model["lr"], + weight_decay=args.weight_decay, + ) + d_model["scheduler"] = optim.lr_scheduler.LinearLR( + d_model["optimizer"], + start_factor=0.001, + end_factor=1.0, + total_iters=30, + verbose=False, + ) + + criterion = nn.MSELoss().to(device) + + print("training loop") + running_losses = {"train": {}, "test": {}} + for k in ["train", "test"]: + running_losses[k] = {d["name"]: [] for d in models} + running_losses[k]["baseline"] = [] + running_losses[k]["baseline_image"] = [] + + saves = [] + + def prep_data(data): + prepared_data = { + "inputs": { + "radio_feature": data["inputs"]["radio_feature"].to(dtype).to(device), + "drone_state": data["inputs"]["drone_state"].to(dtype).to(device), + }, + "labels": data["labels"][..., cols_for_loss].to(dtype).to(device), + } + + assert not data["inputs"]["radio_feature"].isnan().any() + assert not data["inputs"]["radio_feature"].isinf().any() + assert not data["inputs"]["drone_state"].isnan().any() + assert not data["inputs"]["drone_state"].isinf().any() + + return prepared_data + + test_iterator = iter(testloader) + for epoch in range(args.epochs): + for i, data in enumerate(trainloader, 0): + # move to device, do final prep + prepared_data = prep_data(data) + if True: # torch.cuda.amp.autocast(): + for d_model in models: + for p in d_net["model"].parameters(): + if p.isnan().any(): + breakpoint() + loss, losses = model_forward( + d_model, prepared_data, args, "train", update=i, plot=True + ) + if i % args.lr_scheduler_every == args.lr_scheduler_every - 1: + d_model["scheduler"].step() + loss.backward() + running_losses["train"][d_model["name"]].append(losses) + if args.clip > 0: + torch.nn.utils.clip_grad_norm_( + d_net["model"].parameters(), args.clip + ) # clip gradients + d_model["optimizer"].step() + for p in d_net["model"].parameters(): + if p.isnan().any(): + breakpoint() + labels = prepared_data["labels"] + running_losses["train"]["baseline"].append( + { + "baseline": criterion( + labels * 0 + labels.mean(axis=[0, 1], keepdim=True), labels + ).item() + } + ) + if using_images: + running_losses["train"]["baseline_image"].append( + { + "baseline_image": criterion( + label_images * 0 + label_images.mean(), label_images + ).item() + } + ) + + if i % args.print_every == args.print_every - 1: + for idx in np.arange(args.test_mbs): + try: + data = next(test_iterator) + except StopIteration: + test_iterator = iter(testloader) + data = next(test_iterator) + prepared_data = prep_data(data) + with torch.no_grad(): + for d_model in models: + loss, losses = model_forward( + d_model, + prepared_data, + args, + "test", + update=i, + plot=idx == 0, + ) + running_losses["test"][d_model["name"]].append(losses) + labels = prepared_data["labels"] + running_losses["test"]["baseline"].append( + { + "baseline": criterion( + labels * 0 + labels.mean(axis=[0, 1], keepdim=True), + labels, + ).item() + } + ) + if using_images: + running_losses["test"]["baseline_image"].append( + { + "baseline_image": criterion( + label_images * 0 + label_images.mean(), label_images + ).item() + } + ) + + if i == 0 or i % args.save_every == args.save_every - 1: + save(args, running_losses, models, i, args.keep_n_saves) + + if i % args.print_every == args.print_every - 1: + train_baseline_loss = model_to_losses( + running_losses["train"]["baseline"], args.print_every + ) + test_baseline_loss = model_to_losses( + running_losses["test"]["baseline"], args.test_mbs + ) + train_baseline_image_loss = None + test_baseline_image_loss = None + print(f"[{epoch + 1}, {i + 1:5d}]") + if using_images: + train_baseline_image_loss = model_to_losses( + running_losses["train"]["baseline_image"], args.print_every + ) + test_baseline_image_loss = model_to_losses( + running_losses["test"]["baseline_image"], args.test_mbs + ) + print( + f'\tTrain: baseline: {train_baseline_loss["baseline"][-1]:.3f}, baseline_image: {train_baseline_image_loss["baseline_image"][-1]:.3f} , time { (time.time()-start_time)/(i+1) :.3f} / batch' + ) + print( + f'\tTest: baseline: {test_baseline_loss["baseline"][-1]:.3f}, baseline_image: {test_baseline_image_loss["baseline_image"][-1]:.3f} , time { (time.time()-start_time)/(i+1) :.3f} / batch' + ) + else: + print( + f'\tTrain: baseline: {train_baseline_loss["baseline"][-1]:.3f} , time { (time.time()-start_time)/(i+1) :.3f} / batch' + ) + print( + f'\tTest: baseline: {test_baseline_loss["baseline"][-1]:.3f}, time { (time.time()-start_time)/(i+1) :.3f} / batch' + ) + loss_str = "\t" + "\n\t".join( + [ + "%s(%s):(tr)%s,(ts)%s" + % ( + d["name"], + str(d["dead"]), + model_to_loss_str( + running_losses["train"][d["name"]], args.print_every + ), + model_to_loss_str( + running_losses["test"][d["name"]], args.test_mbs + ), + ) + for d in models + ] + ) + print(loss_str) + loss_str = "\n".join( + [ + "\t%s:(tr)\n%s\n\t%s:(ts)\n%s" + % ( + d["name"], + model_to_stats_str( + running_losses["train"][d["name"]], args.print_every + ), + d["name"], + model_to_stats_str( + running_losses["test"][d["name"]], args.test_mbs + ), + ) + for d in models + ] + ) + print("\t\t\t%s" % stats_title()) + print(loss_str) + if i // args.print_every > 2 and i % args.plot_every == args.plot_every - 1: + plot_loss( + running_losses=running_losses["train"], + baseline_loss=train_baseline_loss, + baseline_image_loss=train_baseline_image_loss, + xtick_spacing=args.print_every, + mean_chunk=args.print_every, + output_prefix=args.output_prefix, + fig=loss_figs["train"], + title="Train", + update=i, + ) + plot_loss( + running_losses=running_losses["test"], + baseline_loss=test_baseline_loss, + baseline_image_loss=test_baseline_image_loss, + xtick_spacing=args.print_every, + mean_chunk=args.test_mbs, + output_prefix=args.output_prefix, + fig=loss_figs["test"], + title="Test", + update=i, + ) + if args.plot: + plt.pause(0.5) + + print("Finished Training") # but do we ever really get here? diff --git a/software/model_training_and_inference/13_learn_beamformer.py b/software/model_training_and_inference/13_learn_beamformer.py index 978a1a48..33ae21b6 100644 --- a/software/model_training_and_inference/13_learn_beamformer.py +++ b/software/model_training_and_inference/13_learn_beamformer.py @@ -15,360 +15,477 @@ from torch.utils.data import dataset, random_split from utils.image_utils import labels_to_source_images -from models.models import (SingleSnapshotNet, SnapshotNet, Task1Net, TransformerModel, - UNet, ComplexFFNN, HybridFFNN) -from utils.spf_dataset import SessionsDataset, SessionsDatasetTask2, collate_fn_beamformer - - -torch.set_printoptions(precision=5,sci_mode=False,linewidth=1000) - -output_cols={ # maybe this should get moved to the dataset part... - 'src_pos':[0,1], - 'src_theta':[2], - 'src_dist':[3], - 'det_delta':[4,5], - 'det_theta':[6], - 'det_space':[7] +from models.models import ( + SingleSnapshotNet, + SnapshotNet, + Task1Net, + TransformerModel, + UNet, + ComplexFFNN, + HybridFFNN, +) +from utils.spf_dataset import ( + SessionsDataset, + SessionsDatasetTask2, + collate_fn_beamformer, +) + + +torch.set_printoptions(precision=5, sci_mode=False, linewidth=1000) + +output_cols = { # maybe this should get moved to the dataset part... + "src_pos": [0, 1], + "src_theta": [2], + "src_dist": [3], + "det_delta": [4, 5], + "det_theta": [6], + "det_space": [7], } -input_cols={ - 'det_pos':[0,1], - 'time':[2], - 'space_delta':[3,4], - 'space_theta':[5], - 'space_dist':[6] +input_cols = { + "det_pos": [0, 1], + "time": [2], + "space_delta": [3, 4], + "space_theta": [5], + "space_dist": [6], } -def src_pos_from_radial(inputs,outputs): - det_pos=inputs[:,:,input_cols['det_pos']] - theta=outputs[:,:,output_cols['src_theta']] - dist=outputs[:,:,output_cols['src_dist']] - return torch.stack([torch.cos(theta*np.pi),torch.sin(theta*np.pi)],axis=2)[...,0]*dist+det_pos -#def complexMSE(output, target): +def src_pos_from_radial(inputs, outputs): + det_pos = inputs[:, :, input_cols["det_pos"]] + theta = outputs[:, :, output_cols["src_theta"]] + dist = outputs[:, :, output_cols["src_dist"]] + return ( + torch.stack([torch.cos(theta * np.pi), torch.sin(theta * np.pi)], axis=2)[ + ..., 0 + ] + * dist + + det_pos + ) + + +# def complexMSE(output, target): # loss = torch.mean((output - target).abs()**2) # return loss -def model_forward(d_model,inputs,outputs,args,beamformer): - d_model['optimizer'].zero_grad() - b,s,_=inputs.shape - losses={} - preds=d_model['model'](inputs.reshape(b*s,-1)) - p=0.5 - loss = criterion(preds,outputs.reshape(b*s,-1)) - beam_loss = criterion(preds,beamformer.reshape(b*s,-1).float()) - losses['mse_loss']=loss.item() - losses['beamformer_loss']=beam_loss.item() - if (i%args.print_every)==args.print_every-1: - d_model['fig'].clf() - axs=d_model['fig'].subplots(1,2,sharex=True,sharey=True) - axs[0].plot(preds[0].detach().numpy(),label='prediction') - axs[1].plot(beamformer[0,0].detach().numpy(),label='beamformer') - x=outputs[0,0].argmax().item() - for idx in [0,1]: - axs[idx].axvline(x=x,c='r') - axs[idx].legend() - d_model['fig'].savefig('%s%s_%d.png' % (args.output_prefix,d_model['name'],i)) - d_model['fig'].canvas.draw_idle() - d_model['dead']=np.isclose(preds.var(axis=1).mean().item(),0.0) - return p*loss+(1-p)*beam_loss,losses - -def model_to_losses(running_loss,mean_chunk): - if len(running_loss)==0: - return {} - losses={} - for k in ['baseline','baseline_image','beamformer_loss','mse_loss']: - if k in running_loss[0]: - if '_stats' not in k: - losses[k]=np.log(np.array( [ np.mean([ l[k] for l in running_loss[idx*mean_chunk:(idx+1)*mean_chunk]]) - for idx in range(len(running_loss)//mean_chunk) ])) - else: - losses[k]=[ torch.stack([ l[k] for l in running_loss[idx*mean_chunk:(idx+1)*mean_chunk] ]).mean(axis=0) - for idx in range(len(running_loss)//mean_chunk) ] - return losses - -def model_to_loss_str(running_loss,mean_chunk): - if len(running_loss)==0: - return "" - loss_str=[] - losses=model_to_losses(running_loss,mean_chunk) - for k in ['beamformer_loss','mse_loss']: - if k in losses: - loss_str.append("%s:%0.4f" % (k,losses[k][-1])) - return ",".join(loss_str) + +def model_forward(d_model, inputs, outputs, args, beamformer): + d_model["optimizer"].zero_grad() + b, s, _ = inputs.shape + losses = {} + preds = d_model["model"](inputs.reshape(b * s, -1)) + p = 0.5 + loss = criterion(preds, outputs.reshape(b * s, -1)) + beam_loss = criterion(preds, beamformer.reshape(b * s, -1).float()) + losses["mse_loss"] = loss.item() + losses["beamformer_loss"] = beam_loss.item() + if (i % args.print_every) == args.print_every - 1: + d_model["fig"].clf() + axs = d_model["fig"].subplots(1, 2, sharex=True, sharey=True) + axs[0].plot(preds[0].detach().numpy(), label="prediction") + axs[1].plot(beamformer[0, 0].detach().numpy(), label="beamformer") + x = outputs[0, 0].argmax().item() + for idx in [0, 1]: + axs[idx].axvline(x=x, c="r") + axs[idx].legend() + d_model["fig"].savefig("%s%s_%d.png" % (args.output_prefix, d_model["name"], i)) + d_model["fig"].canvas.draw_idle() + d_model["dead"] = np.isclose(preds.var(axis=1).mean().item(), 0.0) + return p * loss + (1 - p) * beam_loss, losses + + +def model_to_losses(running_loss, mean_chunk): + if len(running_loss) == 0: + return {} + losses = {} + for k in ["baseline", "baseline_image", "beamformer_loss", "mse_loss"]: + if k in running_loss[0]: + if "_stats" not in k: + losses[k] = np.log( + np.array( + [ + np.mean( + [ + l[k] + for l in running_loss[ + idx * mean_chunk : (idx + 1) * mean_chunk + ] + ] + ) + for idx in range(len(running_loss) // mean_chunk) + ] + ) + ) + else: + losses[k] = [ + torch.stack( + [ + l[k] + for l in running_loss[ + idx * mean_chunk : (idx + 1) * mean_chunk + ] + ] + ).mean(axis=0) + for idx in range(len(running_loss) // mean_chunk) + ] + return losses + + +def model_to_loss_str(running_loss, mean_chunk): + if len(running_loss) == 0: + return "" + loss_str = [] + losses = model_to_losses(running_loss, mean_chunk) + for k in ["beamformer_loss", "mse_loss"]: + if k in losses: + loss_str.append("%s:%0.4f" % (k, losses[k][-1])) + return ",".join(loss_str) + def stats_title(): - title_str=[] - for col in args.losses.split(','): - for _ in range(len(output_cols[col])): - title_str.append(col) - return "\t".join(title_str) - - -def model_to_stats_str(running_loss,mean_chunk): - if len(running_loss)==0: - return "" - losses=model_to_losses(running_loss,mean_chunk) - loss_str=[] - for k in ['transformer_stats','single_snapshot_stats']: - if k in losses: - loss_str.append("\t\t%s\t%s" % (k,"\t".join([ "%0.4f" % v.item() for v in losses[k][-1][cols_for_loss]]))) - return "\n".join(loss_str) - -def save(args,running_losses,models,iteration,keep_n_saves): - fn='%ssave_%d.pkl' % (args.output_prefix,iteration) - pickle.dump({ - 'models':models, - 'args':args, - 'running_losses':running_losses},open(fn,'wb')) - saves.append(fn) - while len(saves)>keep_n_saves: - fn=saves.pop(0) - if os.path.exists(fn): - os.remove(fn) - -def plot_loss(running_losses, - baseline_loss, - xtick_spacing, - mean_chunk, - output_prefix, - fig, - title): - fig.clf() - fig.suptitle(title) - ax=fig.subplots(1,2,sharex=True) - xs=np.arange(len(baseline_loss['baseline']))*xtick_spacing - ax[0].plot(xs,baseline_loss['baseline'],label='baseline') - for idx in [0,1]: - ax[idx].set_xlabel("time") - ax[idx].set_ylabel("log loss") - ax[idx].set_title("Loss") - #mn=baseline_loss['baseline'].max()-baseline_loss['baseline'].std() - #print(baseline_loss['baseline'].std()) - for d_model in models: - losses=model_to_losses(running_losses[d_model['name']],mean_chunk) - if 'mse_loss' in losses: - ax[0].plot(xs[2:],losses['mse_loss'][2:],label=d_model['name']) - if 'beamformer_loss' in losses: - ax[1].plot(xs[2:],losses['beamformer_loss'][2:],label=d_model['name']) - #_mn=np.min(losses['beamformer_loss']) - #if _mn0: - torch.nn.utils.clip_grad_norm_(d_net['model'].parameters(), args.clip) # clip gradients - d_model['optimizer'].step() - running_losses['train']['baseline'].append( - {'baseline':criterion(data['beamformer'], data['labels']).item() } ) - - if i%args.print_every==args.print_every-1: - for idx in np.arange(args.test_mbs): - try: - data = next(test_iterator) - except StopIteration: - test_iterator = iter(testloader) - data = next(test_iterator) - data=prep_data(data) - with torch.no_grad(): - for d_model in models: - loss,losses=model_forward(d_model,data['input'],data['labels'],args,data['beamformer']) - running_losses['test'][d_model['name']].append(losses) - running_losses['test']['baseline'].append( - {'baseline':criterion(data['beamformer'], data['labels']).item() } ) - - - if i==0 or i%args.save_every==args.save_every-1: - save(args,running_losses,models,i,args.keep_n_saves) - - if i % args.print_every == args.print_every-1: - - train_baseline_loss=model_to_losses(running_losses['train']['baseline'],args.print_every) - test_baseline_loss=model_to_losses(running_losses['test']['baseline'],args.test_mbs) - print(f'[{epoch + 1}, {i + 1:5d}]') - print(f'\tTrain: baseline: {train_baseline_loss["baseline"][-1]:.3f} , time { (time.time()-start_time)/(i+1) :.3f} / batch' ) - print(f'\tTest: baseline: {test_baseline_loss["baseline"][-1]:.3f}, time { (time.time()-start_time)/(i+1) :.3f} / batch' ) - loss_str="\t"+"\n\t".join( - [ "%s(%s):(tr)%s,(ts)%s" % (d['name'],str(d['dead']), - model_to_loss_str(running_losses['train'][d['name']],args.print_every), - model_to_loss_str(running_losses['test'][d['name']],args.test_mbs) - ) for d in models ]) - print(loss_str) - if i//args.print_every>0: - plot_loss(running_losses=running_losses['train'], - baseline_loss=train_baseline_loss, - xtick_spacing=args.print_every, - mean_chunk=args.print_every, - output_prefix=args.output_prefix, - fig=loss_figs['train'], - title='Train') - plot_loss(running_losses=running_losses['test'], - baseline_loss=test_baseline_loss, - xtick_spacing=args.print_every, - mean_chunk=args.test_mbs, - output_prefix=args.output_prefix, - fig=loss_figs['test'], - title='Test') - if args.plot: - plt.pause(0.5) - - print('Finished Training') # but do we ever really get here? + title_str = [] + for col in args.losses.split(","): + for _ in range(len(output_cols[col])): + title_str.append(col) + return "\t".join(title_str) + + +def model_to_stats_str(running_loss, mean_chunk): + if len(running_loss) == 0: + return "" + losses = model_to_losses(running_loss, mean_chunk) + loss_str = [] + for k in ["transformer_stats", "single_snapshot_stats"]: + if k in losses: + loss_str.append( + "\t\t%s\t%s" + % ( + k, + "\t".join( + ["%0.4f" % v.item() for v in losses[k][-1][cols_for_loss]] + ), + ) + ) + return "\n".join(loss_str) + + +def save(args, running_losses, models, iteration, keep_n_saves): + fn = "%ssave_%d.pkl" % (args.output_prefix, iteration) + pickle.dump( + {"models": models, "args": args, "running_losses": running_losses}, + open(fn, "wb"), + ) + saves.append(fn) + while len(saves) > keep_n_saves: + fn = saves.pop(0) + if os.path.exists(fn): + os.remove(fn) + + +def plot_loss( + running_losses, baseline_loss, xtick_spacing, mean_chunk, output_prefix, fig, title +): + fig.clf() + fig.suptitle(title) + ax = fig.subplots(1, 2, sharex=True) + xs = np.arange(len(baseline_loss["baseline"])) * xtick_spacing + ax[0].plot(xs, baseline_loss["baseline"], label="baseline") + for idx in [0, 1]: + ax[idx].set_xlabel("time") + ax[idx].set_ylabel("log loss") + ax[idx].set_title("Loss") + # mn=baseline_loss['baseline'].max()-baseline_loss['baseline'].std() + # print(baseline_loss['baseline'].std()) + for d_model in models: + losses = model_to_losses(running_losses[d_model["name"]], mean_chunk) + if "mse_loss" in losses: + ax[0].plot(xs[2:], losses["mse_loss"][2:], label=d_model["name"]) + if "beamformer_loss" in losses: + ax[1].plot(xs[2:], losses["beamformer_loss"][2:], label=d_model["name"]) + # _mn=np.min(losses['beamformer_loss']) + # if _mn 0: + torch.nn.utils.clip_grad_norm_( + d_net["model"].parameters(), args.clip + ) # clip gradients + d_model["optimizer"].step() + running_losses["train"]["baseline"].append( + {"baseline": criterion(data["beamformer"], data["labels"]).item()} + ) + + if i % args.print_every == args.print_every - 1: + for idx in np.arange(args.test_mbs): + try: + data = next(test_iterator) + except StopIteration: + test_iterator = iter(testloader) + data = next(test_iterator) + data = prep_data(data) + with torch.no_grad(): + for d_model in models: + loss, losses = model_forward( + d_model, + data["input"], + data["labels"], + args, + data["beamformer"], + ) + running_losses["test"][d_model["name"]].append(losses) + running_losses["test"]["baseline"].append( + { + "baseline": criterion( + data["beamformer"], data["labels"] + ).item() + } + ) + + if i == 0 or i % args.save_every == args.save_every - 1: + save(args, running_losses, models, i, args.keep_n_saves) + + if i % args.print_every == args.print_every - 1: + train_baseline_loss = model_to_losses( + running_losses["train"]["baseline"], args.print_every + ) + test_baseline_loss = model_to_losses( + running_losses["test"]["baseline"], args.test_mbs + ) + print(f"[{epoch + 1}, {i + 1:5d}]") + print( + f'\tTrain: baseline: {train_baseline_loss["baseline"][-1]:.3f} , time { (time.time()-start_time)/(i+1) :.3f} / batch' + ) + print( + f'\tTest: baseline: {test_baseline_loss["baseline"][-1]:.3f}, time { (time.time()-start_time)/(i+1) :.3f} / batch' + ) + loss_str = "\t" + "\n\t".join( + [ + "%s(%s):(tr)%s,(ts)%s" + % ( + d["name"], + str(d["dead"]), + model_to_loss_str( + running_losses["train"][d["name"]], args.print_every + ), + model_to_loss_str( + running_losses["test"][d["name"]], args.test_mbs + ), + ) + for d in models + ] + ) + print(loss_str) + if i // args.print_every > 0: + plot_loss( + running_losses=running_losses["train"], + baseline_loss=train_baseline_loss, + xtick_spacing=args.print_every, + mean_chunk=args.print_every, + output_prefix=args.output_prefix, + fig=loss_figs["train"], + title="Train", + ) + plot_loss( + running_losses=running_losses["test"], + baseline_loss=test_baseline_loss, + xtick_spacing=args.print_every, + mean_chunk=args.test_mbs, + output_prefix=args.output_prefix, + fig=loss_figs["test"], + title="Test", + ) + if args.plot: + plt.pause(0.5) + + print("Finished Training") # but do we ever really get here? diff --git a/software/model_training_and_inference/14_task3_model_training.py b/software/model_training_and_inference/14_task3_model_training.py index 8cacedbe..33f099f2 100644 --- a/software/model_training_and_inference/14_task3_model_training.py +++ b/software/model_training_and_inference/14_task3_model_training.py @@ -14,586 +14,822 @@ from torch.utils.data import dataset, random_split from matplotlib.patches import Ellipse -from models.models import (SingleSnapshotNet, SnapshotNet, Task1Net, TransformerEncOnlyModel, - UNet,TrajectoryNet) -from utils.spf_dataset import SessionsDatasetRealTask2, SessionsDatasetTask2, collate_fn_transformer_filter, output_cols, input_cols +from models.models import ( + SingleSnapshotNet, + SnapshotNet, + Task1Net, + TransformerEncOnlyModel, + UNet, + TrajectoryNet, +) +from utils.spf_dataset import ( + SessionsDatasetRealTask2, + SessionsDatasetTask2, + collate_fn_transformer_filter, + output_cols, + input_cols, +) torch.set_num_threads(8) -torch.set_printoptions(precision=5,sci_mode=False,linewidth=1000) +torch.set_printoptions(precision=5, sci_mode=False, linewidth=1000) def get_rot_mats(theta): - assert(theta.dim()==2 and theta.shape[1]==1) - s = torch.sin(theta) - c = torch.cos(theta) - return torch.cat([c , -s, s, c],axis=1).reshape(theta.shape[0],2,2) + assert theta.dim() == 2 and theta.shape[1] == 1 + s = torch.sin(theta) + c = torch.cos(theta) + return torch.cat([c, -s, s, c], axis=1).reshape(theta.shape[0], 2, 2) + + +# R_i @ vector_i +def rotate_points_by_thetas(points, thetas): + # thetas n x D x D , D is point space dimension, n is number of points + # points n x D + return torch.einsum("ijk,ik->ij", get_rot_mats(thetas), points) -#R_i @ vector_i -def rotate_points_by_thetas(points,thetas): - #thetas n x D x D , D is point space dimension, n is number of points - #points n x D - return torch.einsum('ijk,ik->ij',get_rot_mats(thetas),points) def unpack_mean_cov_angle(x): - return x[:,:2],x[:,2:4],x[:,[4]] - - -#points (n,2) -#means (n,2) -#sigmas (n,2) -#thetas (n,1) -def convert_sigmas(sigmas,min_sigma,max_sigma,ellipse=False): - z=torch.sigmoid(sigmas)*(max_sigma-min_sigma)+min_sigma - if ellipse: - return z - #otherwise make sure its circles - return z.mean(axis=1,keepdim=True).expand_as(z) - -def points_to_nll(points,means,sigmas,thetas,mode='ellipse'): #,min_sigma=0.01,max_sigma=0.3,ellipse=False): - #sigmas=torch.clamp(sigmas.abs(),min=min_sigma,max=None) # TODO clamp hides the gradient? - #sigmas=torch.sigmoid(sigmas)*(max_sigma-min_sigma)+min_sigma #.abs()+min_sigma - if mode=='ellipse': - p=rotate_points_by_thetas(points-means,-thetas)/sigmas # -theta, to undo the rotation, - elif mode=='circle': - p=(points-means)/sigmas.mean(axis=1,keepdim=True) - else: - assert(1==0) - return p.pow(2).sum(axis=1)*0.5 + torch.log(2*torch.pi*sigmas[:,0]*sigmas[:,1]) - -def src_pos_from_radial(inputs,outputs): - det_pos=inputs[:,:,input_cols['det_pos']] - - theta=outputs[:,:,[cols_for_loss.index(output_cols['src_theta'][0])]]*np.pi - dist=outputs[:,:,[cols_for_loss.index(output_cols['src_dist'][0])]] - - theta=theta.float() - dist=dist.float() - det_pos=det_pos.float() - return torch.stack([torch.sin(theta),torch.cos(theta)],axis=2)[...,0]*dist+det_pos - -def model_forward(d_model,data,args,train_test_label,update,plot=True): - batch_size,time_steps,d_drone_state=data['drone_state'].shape - _,_,n_sources,d_emitter_state=data['emitter_position_and_velocity'].shape - - min_sigma=0.01 - max_sigma=0.3 - d_model['optimizer'].zero_grad() - losses={} - preds=d_model['model'](data)#.detach().cpu() - - positions=data['emitter_position_and_velocity'][...,:2].reshape(-1,2) - velocities=data['emitter_position_and_velocity'][...,2:4].reshape(-1,2) - - pos_mean,pos_cov,pos_angle=unpack_mean_cov_angle(preds['trajectory_predictions'][:,:,:,:5].reshape(-1,5)) - vel_mean,vel_cov,vel_angle=unpack_mean_cov_angle(preds['trajectory_predictions'][:,:,:,5:].reshape(-1,5)) - - nll_position_reconstruction_loss=points_to_nll(positions, - pos_mean, - convert_sigmas(pos_cov,min_sigma=min_sigma,max_sigma=max_sigma,ellipse=args.ellipse), - pos_angle,mode=args.point_mode).mean() - nll_velocity_reconstruction_loss=points_to_nll(velocities, - vel_mean, - convert_sigmas(vel_cov,min_sigma=min_sigma,max_sigma=max_sigma,ellipse=args.ellipse), - vel_angle,mode=args.point_mode).mean() - - ss_mean,ss_cov,ss_angle=unpack_mean_cov_angle(preds['single_snapshot_predictions'].reshape(-1,5)) - emitting_positions=data['emitter_position_and_velocity'][data['emitters_broadcasting'][...,0].to(bool)][:,:2] - - nll_ss_position_reconstruction_loss=points_to_nll(emitting_positions, - ss_mean, - convert_sigmas(ss_cov,min_sigma=min_sigma,max_sigma=max_sigma,ellipse=args.ellipse), - ss_angle,mode=args.point_mode).mean() - if plot and (update%args.plot_every)==args.plot_every-1: - t=time_steps - color=['g', 'b','orange', 'y'] - d_model['fig'].clf() - axs=d_model['fig'].subplots(1,3,sharex=True,sharey=True,subplot_kw=dict(box_aspect=1)) - axs[0].set_xlim([-1,1]) - axs[0].set_ylim([-1,1]) - #axs[0].scatter(_l[0,:,src_pos_idxs[0]],_l[0,:,src_pos_idxs[1]],label='source positions',c='r') - _emitting_positions=emitting_positions[:t].cpu() - axs[0].scatter(data['drone_state'][0,:t,0].cpu(), - data['drone_state'][0,:t,1].cpu(),label='detector positions',s=2) - #emitters - axs[0].scatter(_emitting_positions[:,0],_emitting_positions[:,1],label='emitting positions',c='r',alpha=0.1,s=20,facecolor='none') - #all source trajectories - for source_idx in range(n_sources): - _positions=data['emitter_position_and_velocity'][0,:,source_idx,:2].cpu().detach().numpy() - axs[0].scatter( - _positions[:t,0], - _positions[:t,1], - label='emitter %d' % (source_idx+1),c=color[source_idx%len(color)],alpha=0.3,s=1) - - - axs[0].set_title("Ground truth") - axs[1].set_title("Single snapshot predictions") - axs[2].set_title("Trajectory predictions") - _ss_mean=ss_mean[:t].cpu().detach().numpy() - _ss_cov=convert_sigmas(ss_cov[:t],min_sigma=min_sigma,max_sigma=max_sigma,ellipse=args.ellipse).cpu().detach().numpy() - _ss_angle=ss_angle[:t].cpu().detach().numpy() - axs[1].scatter(_ss_mean[:,0],_ss_mean[:,1],label='pred means',c='r',alpha=0.3,s=2) - print("PLOT",_ss_cov[:t].mean(),_ss_cov[:t].max()) - for idx in torch.arange(0,t,5): - ellipse = Ellipse((_ss_mean[idx,0], _ss_mean[idx,1]), - width=_ss_cov[idx,0]*3, - height=_ss_cov[idx,1]*3, - facecolor='none',edgecolor='red',alpha=0.1, - angle=_ss_angle[idx]*(360.0/(2*torch.pi) if args.ellipse else 0) - ) - axs[1].add_patch(ellipse) - - axs[1].set_title("Single snapshot predictions") - - _pred_trajectory=preds['trajectory_predictions']#.cpu().detach().numpy() - for source_idx in range(n_sources): - trajectory_mean,trajectory_cov,trajectory_angle=unpack_mean_cov_angle(_pred_trajectory[0,:t,source_idx,:5].reshape(-1,5)) - trajectory_mean=trajectory_mean.cpu() - trajectory_cov=trajectory_cov.cpu() - trajectory_angle=trajectory_angle.cpu() - axs[2].scatter( - trajectory_mean[:,0].detach().numpy(), - trajectory_mean[:,1].detach().numpy(), - label='trajectory emitter %d' % (source_idx+1),s=2,color=color[source_idx%len(color)]) - - trajectory_cov=convert_sigmas(trajectory_cov,min_sigma=min_sigma,max_sigma=max_sigma,ellipse=args.ellipse).cpu().detach().numpy() - for idx in torch.arange(0,t,5): - ellipse = Ellipse((trajectory_mean[idx,0], trajectory_mean[idx,1]), - width=trajectory_cov[idx,0]*3, - height=trajectory_cov[idx,1]*3, - facecolor='none', - edgecolor=color[source_idx%len(color)], - alpha=0.1, - angle=-trajectory_angle[idx]*(360.0/(2*torch.pi) if args.ellipse else 0) + return x[:, :2], x[:, 2:4], x[:, [4]] + + +# points (n,2) +# means (n,2) +# sigmas (n,2) +# thetas (n,1) +def convert_sigmas(sigmas, min_sigma, max_sigma, ellipse=False): + z = torch.sigmoid(sigmas) * (max_sigma - min_sigma) + min_sigma + if ellipse: + return z + # otherwise make sure its circles + return z.mean(axis=1, keepdim=True).expand_as(z) + + +def points_to_nll( + points, means, sigmas, thetas, mode="ellipse" +): # ,min_sigma=0.01,max_sigma=0.3,ellipse=False): + # sigmas=torch.clamp(sigmas.abs(),min=min_sigma,max=None) # TODO clamp hides the gradient? + # sigmas=torch.sigmoid(sigmas)*(max_sigma-min_sigma)+min_sigma #.abs()+min_sigma + if mode == "ellipse": + p = ( + rotate_points_by_thetas(points - means, -thetas) / sigmas + ) # -theta, to undo the rotation, + elif mode == "circle": + p = (points - means) / sigmas.mean(axis=1, keepdim=True) + else: + assert 1 == 0 + return p.pow(2).sum(axis=1) * 0.5 + torch.log( + 2 * torch.pi * sigmas[:, 0] * sigmas[:, 1] + ) + + +def src_pos_from_radial(inputs, outputs): + det_pos = inputs[:, :, input_cols["det_pos"]] + + theta = outputs[:, :, [cols_for_loss.index(output_cols["src_theta"][0])]] * np.pi + dist = outputs[:, :, [cols_for_loss.index(output_cols["src_dist"][0])]] + + theta = theta.float() + dist = dist.float() + det_pos = det_pos.float() + return ( + torch.stack([torch.sin(theta), torch.cos(theta)], axis=2)[..., 0] * dist + + det_pos + ) + + +def model_forward(d_model, data, args, train_test_label, update, plot=True): + batch_size, time_steps, d_drone_state = data["drone_state"].shape + _, _, n_sources, d_emitter_state = data["emitter_position_and_velocity"].shape + + min_sigma = 0.01 + max_sigma = 0.3 + d_model["optimizer"].zero_grad() + losses = {} + preds = d_model["model"](data) # .detach().cpu() + + positions = data["emitter_position_and_velocity"][..., :2].reshape(-1, 2) + velocities = data["emitter_position_and_velocity"][..., 2:4].reshape(-1, 2) + + pos_mean, pos_cov, pos_angle = unpack_mean_cov_angle( + preds["trajectory_predictions"][:, :, :, :5].reshape(-1, 5) + ) + vel_mean, vel_cov, vel_angle = unpack_mean_cov_angle( + preds["trajectory_predictions"][:, :, :, 5:].reshape(-1, 5) + ) + + nll_position_reconstruction_loss = points_to_nll( + positions, + pos_mean, + convert_sigmas( + pos_cov, min_sigma=min_sigma, max_sigma=max_sigma, ellipse=args.ellipse + ), + pos_angle, + mode=args.point_mode, + ).mean() + nll_velocity_reconstruction_loss = points_to_nll( + velocities, + vel_mean, + convert_sigmas( + vel_cov, min_sigma=min_sigma, max_sigma=max_sigma, ellipse=args.ellipse + ), + vel_angle, + mode=args.point_mode, + ).mean() + + ss_mean, ss_cov, ss_angle = unpack_mean_cov_angle( + preds["single_snapshot_predictions"].reshape(-1, 5) + ) + emitting_positions = data["emitter_position_and_velocity"][ + data["emitters_broadcasting"][..., 0].to(bool) + ][:, :2] + + nll_ss_position_reconstruction_loss = points_to_nll( + emitting_positions, + ss_mean, + convert_sigmas( + ss_cov, min_sigma=min_sigma, max_sigma=max_sigma, ellipse=args.ellipse + ), + ss_angle, + mode=args.point_mode, + ).mean() + if plot and (update % args.plot_every) == args.plot_every - 1: + t = time_steps + color = ["g", "b", "orange", "y"] + d_model["fig"].clf() + axs = d_model["fig"].subplots( + 1, 3, sharex=True, sharey=True, subplot_kw=dict(box_aspect=1) + ) + axs[0].set_xlim([-1, 1]) + axs[0].set_ylim([-1, 1]) + # axs[0].scatter(_l[0,:,src_pos_idxs[0]],_l[0,:,src_pos_idxs[1]],label='source positions',c='r') + _emitting_positions = emitting_positions[:t].cpu() + axs[0].scatter( + data["drone_state"][0, :t, 0].cpu(), + data["drone_state"][0, :t, 1].cpu(), + label="detector positions", + s=2, + ) + # emitters + axs[0].scatter( + _emitting_positions[:, 0], + _emitting_positions[:, 1], + label="emitting positions", + c="r", + alpha=0.1, + s=20, + facecolor="none", + ) + # all source trajectories + for source_idx in range(n_sources): + _positions = ( + data["emitter_position_and_velocity"][0, :, source_idx, :2] + .cpu() + .detach() + .numpy() + ) + axs[0].scatter( + _positions[:t, 0], + _positions[:t, 1], + label="emitter %d" % (source_idx + 1), + c=color[source_idx % len(color)], + alpha=0.3, + s=1, ) - axs[2].add_patch(ellipse) - for idx in [0,1,2]: - axs[idx].legend() - axs[idx].set_xlabel("X") - axs[idx].set_ylabel("Y") - # - # axs[2].scatter(_l[0,:,src_pos_idxs[0]],_l[0,:,src_pos_idxs[1]],label='real positions',c='b',alpha=0.1,s=7) - # axs[2].scatter(_p[0,:,src_pos_idxs[0]],_p[0,:,src_pos_idxs[1]],label='predicted positions',c='r',alpha=0.3,s=7) - d_model['fig'].tight_layout() - d_model['fig'].canvas.draw_idle() - d_model['fig'].savefig('%s%s_%d_%s.png' % (args.output_prefix,d_model['name'],update,train_test_label)) - - losses={ - 'nll_position_reconstruction_loss':nll_position_reconstruction_loss.item(), - 'nll_velocity_reconstruction_loss':nll_velocity_reconstruction_loss.item(), - 'nll_ss_position_reconstruction_loss':nll_ss_position_reconstruction_loss.item() - } - #loss=nll_position_reconstruction_loss+nll_velocity_reconstruction_loss+nll_ss_position_reconstruction_loss - lm=torch.tensor([args.loss_single,args.loss_trec,args.loss_tvel]) - lm/=lm.sum() - loss=lm[0]*nll_ss_position_reconstruction_loss+lm[1]*nll_position_reconstruction_loss+lm[2]*nll_velocity_reconstruction_loss - - if args.l2!=0.0: - l2s=d_model['model'].l2() - for k in l2s: - loss+=args.l2*l2s[k] - if args.embedding_warmup>update: - loss=nll_ss_position_reconstruction_loss - return loss,losses - -def model_to_losses(running_loss,mean_chunk): - if len(running_loss)==0: - return {} - losses={} - for k in ['baseline','nll_position_reconstruction_loss','nll_velocity_reconstruction_loss','nll_ss_position_reconstruction_loss']: - if k in running_loss[0]: - if '_stats' not in k: - #losses[k]=np.log(np.array( [ np.mean([ l[k] for l in running_loss[idx*mean_chunk:(idx+1)*mean_chunk]]) - # for idx in range(len(running_loss)//mean_chunk) ])) - losses[k]=np.array( [ np.mean([ l[k] for l in running_loss[idx*mean_chunk:(idx+1)*mean_chunk]]) - for idx in range(len(running_loss)//mean_chunk) ]) - else: - losses[k]=[ torch.stack([ l[k] for l in running_loss[idx*mean_chunk:(idx+1)*mean_chunk] ]).mean(axis=0) - for idx in range(len(running_loss)//mean_chunk) ] - return losses - -def model_to_loss_str(running_loss,mean_chunk): - if len(running_loss)==0: - return "" - loss_str=[] - losses=model_to_losses(running_loss,mean_chunk) - for k in ['nll_position_reconstruction_loss','nll_velocity_reconstruction_loss','nll_ss_position_reconstruction_loss']: - if k in losses: - loss_str.append("%s:%0.4f" % (k,losses[k][-1])) - return ",".join(loss_str) + + axs[0].set_title("Ground truth") + axs[1].set_title("Single snapshot predictions") + axs[2].set_title("Trajectory predictions") + _ss_mean = ss_mean[:t].cpu().detach().numpy() + _ss_cov = ( + convert_sigmas( + ss_cov[:t], + min_sigma=min_sigma, + max_sigma=max_sigma, + ellipse=args.ellipse, + ) + .cpu() + .detach() + .numpy() + ) + _ss_angle = ss_angle[:t].cpu().detach().numpy() + axs[1].scatter( + _ss_mean[:, 0], _ss_mean[:, 1], label="pred means", c="r", alpha=0.3, s=2 + ) + print("PLOT", _ss_cov[:t].mean(), _ss_cov[:t].max()) + for idx in torch.arange(0, t, 5): + ellipse = Ellipse( + (_ss_mean[idx, 0], _ss_mean[idx, 1]), + width=_ss_cov[idx, 0] * 3, + height=_ss_cov[idx, 1] * 3, + facecolor="none", + edgecolor="red", + alpha=0.1, + angle=_ss_angle[idx] * (360.0 / (2 * torch.pi) if args.ellipse else 0), + ) + axs[1].add_patch(ellipse) + + axs[1].set_title("Single snapshot predictions") + + _pred_trajectory = preds["trajectory_predictions"] # .cpu().detach().numpy() + for source_idx in range(n_sources): + trajectory_mean, trajectory_cov, trajectory_angle = unpack_mean_cov_angle( + _pred_trajectory[0, :t, source_idx, :5].reshape(-1, 5) + ) + trajectory_mean = trajectory_mean.cpu() + trajectory_cov = trajectory_cov.cpu() + trajectory_angle = trajectory_angle.cpu() + axs[2].scatter( + trajectory_mean[:, 0].detach().numpy(), + trajectory_mean[:, 1].detach().numpy(), + label="trajectory emitter %d" % (source_idx + 1), + s=2, + color=color[source_idx % len(color)], + ) + + trajectory_cov = ( + convert_sigmas( + trajectory_cov, + min_sigma=min_sigma, + max_sigma=max_sigma, + ellipse=args.ellipse, + ) + .cpu() + .detach() + .numpy() + ) + for idx in torch.arange(0, t, 5): + ellipse = Ellipse( + (trajectory_mean[idx, 0], trajectory_mean[idx, 1]), + width=trajectory_cov[idx, 0] * 3, + height=trajectory_cov[idx, 1] * 3, + facecolor="none", + edgecolor=color[source_idx % len(color)], + alpha=0.1, + angle=-trajectory_angle[idx] + * (360.0 / (2 * torch.pi) if args.ellipse else 0), + ) + axs[2].add_patch(ellipse) + for idx in [0, 1, 2]: + axs[idx].legend() + axs[idx].set_xlabel("X") + axs[idx].set_ylabel("Y") + # + # axs[2].scatter(_l[0,:,src_pos_idxs[0]],_l[0,:,src_pos_idxs[1]],label='real positions',c='b',alpha=0.1,s=7) + # axs[2].scatter(_p[0,:,src_pos_idxs[0]],_p[0,:,src_pos_idxs[1]],label='predicted positions',c='r',alpha=0.3,s=7) + d_model["fig"].tight_layout() + d_model["fig"].canvas.draw_idle() + d_model["fig"].savefig( + "%s%s_%d_%s.png" + % (args.output_prefix, d_model["name"], update, train_test_label) + ) + + losses = { + "nll_position_reconstruction_loss": nll_position_reconstruction_loss.item(), + "nll_velocity_reconstruction_loss": nll_velocity_reconstruction_loss.item(), + "nll_ss_position_reconstruction_loss": nll_ss_position_reconstruction_loss.item(), + } + # loss=nll_position_reconstruction_loss+nll_velocity_reconstruction_loss+nll_ss_position_reconstruction_loss + lm = torch.tensor([args.loss_single, args.loss_trec, args.loss_tvel]) + lm /= lm.sum() + loss = ( + lm[0] * nll_ss_position_reconstruction_loss + + lm[1] * nll_position_reconstruction_loss + + lm[2] * nll_velocity_reconstruction_loss + ) + + if args.l2 != 0.0: + l2s = d_model["model"].l2() + for k in l2s: + loss += args.l2 * l2s[k] + if args.embedding_warmup > update: + loss = nll_ss_position_reconstruction_loss + return loss, losses + + +def model_to_losses(running_loss, mean_chunk): + if len(running_loss) == 0: + return {} + losses = {} + for k in [ + "baseline", + "nll_position_reconstruction_loss", + "nll_velocity_reconstruction_loss", + "nll_ss_position_reconstruction_loss", + ]: + if k in running_loss[0]: + if "_stats" not in k: + # losses[k]=np.log(np.array( [ np.mean([ l[k] for l in running_loss[idx*mean_chunk:(idx+1)*mean_chunk]]) + # for idx in range(len(running_loss)//mean_chunk) ])) + losses[k] = np.array( + [ + np.mean( + [ + l[k] + for l in running_loss[ + idx * mean_chunk : (idx + 1) * mean_chunk + ] + ] + ) + for idx in range(len(running_loss) // mean_chunk) + ] + ) + else: + losses[k] = [ + torch.stack( + [ + l[k] + for l in running_loss[ + idx * mean_chunk : (idx + 1) * mean_chunk + ] + ] + ).mean(axis=0) + for idx in range(len(running_loss) // mean_chunk) + ] + return losses + + +def model_to_loss_str(running_loss, mean_chunk): + if len(running_loss) == 0: + return "" + loss_str = [] + losses = model_to_losses(running_loss, mean_chunk) + for k in [ + "nll_position_reconstruction_loss", + "nll_velocity_reconstruction_loss", + "nll_ss_position_reconstruction_loss", + ]: + if k in losses: + loss_str.append("%s:%0.4f" % (k, losses[k][-1])) + return ",".join(loss_str) + def stats_title(): - title_str=[] - for col in args.losses.split(','): - for _ in range(len(output_cols[col])): - title_str.append(col) - return "\t".join(title_str) - - -def model_to_stats_str(running_loss,mean_chunk): - if len(running_loss)==0: - return "" - losses=model_to_losses(running_loss,mean_chunk) - loss_str=[] - for k in ['transformer_stats','single_snapshot_stats']: - if k in losses: - loss_str.append("\t\t%s\t%s" % (k,"\t".join([ "%0.4f" % v.item() for v in losses[k][-1]]))) - return "\n".join(loss_str) - -def save(args,running_losses,models,iteration,keep_n_saves): - fn='%ssave_%d.pkl' % (args.output_prefix,iteration) - pickle.dump({ - 'models':models, - 'args':args, - 'running_losses':running_losses},open(fn,'wb')) - saves.append(fn) - while len(saves)>keep_n_saves: - fn=saves.pop(0) - if os.path.exists(fn): - os.remove(fn) - -def plot_loss(running_losses, + title_str = [] + for col in args.losses.split(","): + for _ in range(len(output_cols[col])): + title_str.append(col) + return "\t".join(title_str) + + +def model_to_stats_str(running_loss, mean_chunk): + if len(running_loss) == 0: + return "" + losses = model_to_losses(running_loss, mean_chunk) + loss_str = [] + for k in ["transformer_stats", "single_snapshot_stats"]: + if k in losses: + loss_str.append( + "\t\t%s\t%s" + % (k, "\t".join(["%0.4f" % v.item() for v in losses[k][-1]])) + ) + return "\n".join(loss_str) + + +def save(args, running_losses, models, iteration, keep_n_saves): + fn = "%ssave_%d.pkl" % (args.output_prefix, iteration) + pickle.dump( + {"models": models, "args": args, "running_losses": running_losses}, + open(fn, "wb"), + ) + saves.append(fn) + while len(saves) > keep_n_saves: + fn = saves.pop(0) + if os.path.exists(fn): + os.remove(fn) + + +def plot_loss( + running_losses, baseline_loss, xtick_spacing, mean_chunk, output_prefix, fig, title, - update): - fig.clf() - fig.suptitle(title) - axs=fig.subplots(1,4,sharex=True) - axs[1].sharex(axs[0]) - axs[2].sharex(axs[0]) - xs=np.arange(len(baseline_loss['baseline']))*xtick_spacing - start_idx=min(len(baseline_loss['baseline'])-1,int(len(baseline_loss['baseline'])*0.1)) - for i in range(3): - #axs[i].plot(xs,baseline_loss['baseline'],label='baseline') - axs[i].set_xlabel("time") - axs[i].set_ylabel("log loss") - #for k in ['nll_position_reconstruction_loss','nll_velocity_reconstruction_loss','nll_ss_position_reconstruction_loss']: - axs[0].set_title("nll_ss_position_reconstruction_loss") - axs[1].set_title("nll_position_reconstruction_loss") - axs[2].set_title("nll_velocity_reconstruction_loss") - #axs[3].set_title("Image loss") - for d_model in models: - losses=model_to_losses(running_losses[d_model['name']],mean_chunk) - if 'nll_ss_position_reconstruction_loss' in losses: - axs[0].plot( - xs[start_idx:], - losses['nll_ss_position_reconstruction_loss'][start_idx:],label=d_model['name']) - if 'nll_position_reconstruction_loss' in losses: - axs[1].plot( - xs[start_idx:], - losses['nll_position_reconstruction_loss'][start_idx:],label=d_model['name']) - if 'nll_velocity_reconstruction_loss' in losses: - axs[2].plot( - xs[start_idx:], - losses['nll_velocity_reconstruction_loss'][start_idx:],label=d_model['name']) - for i in range(4): - axs[i].legend() - fig.tight_layout() - fig.savefig('%sloss_%s_%d.png' % (output_prefix,title,update)) - fig.canvas.draw_idle() - -if __name__=='__main__': - parser = argparse.ArgumentParser() - parser.add_argument('--device', type=str, required=False, default='cpu') - parser.add_argument('--embedding-warmup', type=int, required=False, default=4096) - parser.add_argument('--snapshots-per-sample', type=int, required=False, default=[1,4,8], nargs="+") - parser.add_argument('--n-layers', type=int, required=False, default=[4], nargs="+") - parser.add_argument('--print-every', type=int, required=False, default=100) - parser.add_argument('--lr-scheduler-every', type=int, required=False, default=256) - parser.add_argument('--step-size', type=int, required=False, default=32) - parser.add_argument('--plot-every', type=int, required=False, default=1024) - parser.add_argument('--save-every', type=int, required=False, default=1000) - parser.add_argument('--loss-single', type=float, required=False, default=3.0) - parser.add_argument('--loss-trec', type=float, required=False, default=3.0) - parser.add_argument('--loss-tvel', type=float, required=False, default=0.0) - parser.add_argument('--l2', type=float, required=False, default=0.0001) - parser.add_argument('--ellipse', type=bool, required=False, default=True) - parser.add_argument('--point-mode', type=str, required=False, default='ellipse') - parser.add_argument('--real-data', action='store_true') - parser.add_argument('--test-mbs', type=int, required=False, default=8) - parser.add_argument('--beam-former-spacing', type=int, required=False, default=256+1) - parser.add_argument('--output-prefix', type=str, required=False, default='model_out') - parser.add_argument('--test-fraction', type=float, required=False, default=0.2) - parser.add_argument('--weight-decay', type=float, required=False, default=0.0) - parser.add_argument('--transformer-loss-balance', type=float, required=False, default=0.1) - parser.add_argument('--type', type=str, required=False, default=32) - parser.add_argument('--seed', type=int, required=False, default=0) - parser.add_argument('--keep-n-saves', type=int, required=False, default=2) - parser.add_argument('--epochs', type=int, required=False, default=20000) - parser.add_argument('--positional-encoding-len', type=int, required=False, default=0) - parser.add_argument('--mb', type=int, required=False, default=64) - parser.add_argument('--workers', type=int, required=False, default=4) - parser.add_argument('--dataset', type=str, required=False, default='./sessions-default') - parser.add_argument('--lr-image', type=float, required=False, default=0.05) - parser.add_argument('--lr-direct', type=float, required=False, default=0.01) - parser.add_argument('--lr-transformer', type=float, required=False, default=0.00001) - parser.add_argument('--plot', type=bool, required=False, default=False) - parser.add_argument('--transformer-input', type=str, required=False, default=['drone_state','embedding','single_snapshot_pred'],nargs="+") - parser.add_argument('--transformer-dmodel', type=int, required=False, default=64) #512) - parser.add_argument('--ssn-dhid', type=int, required=False, default=16) #64) - parser.add_argument('--ssn-nlayers', type=int, required=False, default=3) #8) - parser.add_argument('--tj-dembed', type=int, required=False, default=32) #256) - parser.add_argument('--obv-dembed', type=int, required=False, default=31) #256) - parser.add_argument('--clip', type=float, required=False, default=0.5) - parser.add_argument('--losses', type=str, required=False, default="src_pos,src_theta,src_dist") #,src_theta,src_dist,det_delta,det_theta,det_space") - args = parser.parse_args() - - dtype=torch.float32 - if args.type=='16': - dtype=torch.float16 - - if args.plot==False: - import matplotlib - matplotlib.use('Agg') - import matplotlib.pyplot as plt - - - torch.manual_seed(args.seed) - random.seed(args.seed) - np.random.seed(args.seed) - - - #lets see if output dir exists , if not make it - basename=os.path.basename(args.output_prefix) - if not os.path.exists(basename): - os.makedirs(basename) - - - cols_for_loss=[] - for k in output_cols: - if k in args.losses.split(','): - cols_for_loss+=output_cols[k] - - device=torch.device(args.device) - print("init dataset",args.real_data) - if args.real_data: - snapshots_per_sample=max(args.snapshots_per_sample) - ds=SessionsDatasetRealTask2( - args.dataset, - snapshots_in_sample=snapshots_per_sample, - step_size=args.step_size + update, +): + fig.clf() + fig.suptitle(title) + axs = fig.subplots(1, 4, sharex=True) + axs[1].sharex(axs[0]) + axs[2].sharex(axs[0]) + xs = np.arange(len(baseline_loss["baseline"])) * xtick_spacing + start_idx = min( + len(baseline_loss["baseline"]) - 1, int(len(baseline_loss["baseline"]) * 0.1) ) - else: - ds=SessionsDatasetTask2(args.dataset,snapshots_in_sample=max(args.snapshots_per_sample)) - #ds_test=SessionsDatasetTask2(args.test_dataset,snapshots_in_sample=max(args.snapshots_per_sample)) - train_size=int(len(ds)*args.test_fraction) - test_size=len(ds)-train_size - - #need to generate separate files for this not to train test leak - #ds_train, ds_test = random_split(ds, [1-args.test_fraction, args.test_fraction]) - print("NO SHUFFLING! make sure datasets are shuffled!!!") - ds_train = torch.utils.data.Subset(ds, np.arange(train_size)) - ds_test = torch.utils.data.Subset(ds, np.arange(train_size, train_size + test_size)) - print("init dataloader") - trainloader = torch.utils.data.DataLoader( - ds_train, - batch_size=args.mb, - shuffle=True, - num_workers=args.workers, - collate_fn=collate_fn_transformer_filter) - testloader = torch.utils.data.DataLoader( - ds_test, - batch_size=args.mb, - shuffle=True, - num_workers=args.workers, - collate_fn=collate_fn_transformer_filter) - - print("init network") - models=[] - - if True: - for n_layers in args.n_layers: - models.append({ - 'name':'TrajectoryNet_l%d' % n_layers, - 'model':TrajectoryNet( - d_drone_state=3+4, - d_radio_feature=args.beam_former_spacing+1, # add one for mean power normalization - d_detector_observation_embedding=args.obv_dembed, - d_trajectory_embedding=args.tj_dembed, - trajectory_prediction_n_layers=8, - d_trajectory_prediction_output=(2+2+1)+(2+2+1), - d_model=args.transformer_dmodel, - n_heads=8, - d_hid=64, - n_layers=n_layers, - n_outputs=8, - ssn_d_hid=args.ssn_dhid, - ssn_n_layers=args.ssn_nlayers, - #ssn_d_output=5 - ), - 'fig':plt.figure(figsize=(18,4)), - 'dead':False, - 'lr':args.lr_transformer, - }) - - - if False: - for n_layers in [2,4,8,16,32]:#,32,64]: #,32,64]: - for snapshots_per_sample in args.snapshots_per_sample: - if snapshots_per_sample==1: - continue - models.append( - { - 'name':'%d snapshots (l%d)' % (snapshots_per_sample,n_layers), - 'model':SnapshotNet( - snapshots_per_sample, - n_layers=n_layers, - d_model=args.transformer_dmodel, - n_outputs=len(cols_for_loss), - ssn_n_outputs=len(cols_for_loss), - dropout=0.0, - positional_encoding_len=args.positional_encoding_len, - tformer_input=args.transformer_input), - 'snapshots_per_sample':snapshots_per_sample, - 'images':False, - 'lr':args.lr_transformer, - 'images':False,'fig':plt.figure(figsize=(18,4)), - 'dead':False, - } + for i in range(3): + # axs[i].plot(xs,baseline_loss['baseline'],label='baseline') + axs[i].set_xlabel("time") + axs[i].set_ylabel("log loss") + # for k in ['nll_position_reconstruction_loss','nll_velocity_reconstruction_loss','nll_ss_position_reconstruction_loss']: + axs[0].set_title("nll_ss_position_reconstruction_loss") + axs[1].set_title("nll_position_reconstruction_loss") + axs[2].set_title("nll_velocity_reconstruction_loss") + # axs[3].set_title("Image loss") + for d_model in models: + losses = model_to_losses(running_losses[d_model["name"]], mean_chunk) + if "nll_ss_position_reconstruction_loss" in losses: + axs[0].plot( + xs[start_idx:], + losses["nll_ss_position_reconstruction_loss"][start_idx:], + label=d_model["name"], + ) + if "nll_position_reconstruction_loss" in losses: + axs[1].plot( + xs[start_idx:], + losses["nll_position_reconstruction_loss"][start_idx:], + label=d_model["name"], + ) + if "nll_velocity_reconstruction_loss" in losses: + axs[2].plot( + xs[start_idx:], + losses["nll_velocity_reconstruction_loss"][start_idx:], + label=d_model["name"], + ) + for i in range(4): + axs[i].legend() + fig.tight_layout() + fig.savefig("%sloss_%s_%d.png" % (output_prefix, title, update)) + fig.canvas.draw_idle() + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--device", type=str, required=False, default="cpu") + parser.add_argument("--embedding-warmup", type=int, required=False, default=4096) + parser.add_argument( + "--snapshots-per-sample", type=int, required=False, default=[1, 4, 8], nargs="+" + ) + parser.add_argument("--n-layers", type=int, required=False, default=[4], nargs="+") + parser.add_argument("--print-every", type=int, required=False, default=100) + parser.add_argument("--lr-scheduler-every", type=int, required=False, default=256) + parser.add_argument("--step-size", type=int, required=False, default=32) + parser.add_argument("--plot-every", type=int, required=False, default=1024) + parser.add_argument("--save-every", type=int, required=False, default=1000) + parser.add_argument("--loss-single", type=float, required=False, default=3.0) + parser.add_argument("--loss-trec", type=float, required=False, default=3.0) + parser.add_argument("--loss-tvel", type=float, required=False, default=0.0) + parser.add_argument("--l2", type=float, required=False, default=0.0001) + parser.add_argument("--ellipse", type=bool, required=False, default=True) + parser.add_argument("--point-mode", type=str, required=False, default="ellipse") + parser.add_argument("--real-data", action="store_true") + parser.add_argument("--test-mbs", type=int, required=False, default=8) + parser.add_argument( + "--beam-former-spacing", type=int, required=False, default=256 + 1 + ) + parser.add_argument( + "--output-prefix", type=str, required=False, default="model_out" + ) + parser.add_argument("--test-fraction", type=float, required=False, default=0.2) + parser.add_argument("--weight-decay", type=float, required=False, default=0.0) + parser.add_argument( + "--transformer-loss-balance", type=float, required=False, default=0.1 + ) + parser.add_argument("--type", type=str, required=False, default=32) + parser.add_argument("--seed", type=int, required=False, default=0) + parser.add_argument("--keep-n-saves", type=int, required=False, default=2) + parser.add_argument("--epochs", type=int, required=False, default=20000) + parser.add_argument( + "--positional-encoding-len", type=int, required=False, default=0 + ) + parser.add_argument("--mb", type=int, required=False, default=64) + parser.add_argument("--workers", type=int, required=False, default=4) + parser.add_argument( + "--dataset", type=str, required=False, default="./sessions-default" + ) + parser.add_argument("--lr-image", type=float, required=False, default=0.05) + parser.add_argument("--lr-direct", type=float, required=False, default=0.01) + parser.add_argument("--lr-transformer", type=float, required=False, default=0.00001) + parser.add_argument("--plot", type=bool, required=False, default=False) + parser.add_argument( + "--transformer-input", + type=str, + required=False, + default=["drone_state", "embedding", "single_snapshot_pred"], + nargs="+", + ) + parser.add_argument( + "--transformer-dmodel", type=int, required=False, default=64 + ) # 512) + parser.add_argument("--ssn-dhid", type=int, required=False, default=16) # 64) + parser.add_argument("--ssn-nlayers", type=int, required=False, default=3) # 8) + parser.add_argument("--tj-dembed", type=int, required=False, default=32) # 256) + parser.add_argument("--obv-dembed", type=int, required=False, default=31) # 256) + parser.add_argument("--clip", type=float, required=False, default=0.5) + parser.add_argument( + "--losses", type=str, required=False, default="src_pos,src_theta,src_dist" + ) # ,src_theta,src_dist,det_delta,det_theta,det_space") + args = parser.parse_args() + + dtype = torch.float32 + if args.type == "16": + dtype = torch.float16 + + if args.plot == False: + import matplotlib + + matplotlib.use("Agg") + import matplotlib.pyplot as plt + + torch.manual_seed(args.seed) + random.seed(args.seed) + np.random.seed(args.seed) + + # lets see if output dir exists , if not make it + basename = os.path.basename(args.output_prefix) + if not os.path.exists(basename): + os.makedirs(basename) + + cols_for_loss = [] + for k in output_cols: + if k in args.losses.split(","): + cols_for_loss += output_cols[k] + + device = torch.device(args.device) + print("init dataset", args.real_data) + if args.real_data: + snapshots_per_sample = max(args.snapshots_per_sample) + ds = SessionsDatasetRealTask2( + args.dataset, + snapshots_in_sample=snapshots_per_sample, + step_size=args.step_size, + ) + else: + ds = SessionsDatasetTask2( + args.dataset, snapshots_in_sample=max(args.snapshots_per_sample) ) + # ds_test=SessionsDatasetTask2(args.test_dataset,snapshots_in_sample=max(args.snapshots_per_sample)) + train_size = int(len(ds) * args.test_fraction) + test_size = len(ds) - train_size + + # need to generate separate files for this not to train test leak + # ds_train, ds_test = random_split(ds, [1-args.test_fraction, args.test_fraction]) + print("NO SHUFFLING! make sure datasets are shuffled!!!") + ds_train = torch.utils.data.Subset(ds, np.arange(train_size)) + ds_test = torch.utils.data.Subset(ds, np.arange(train_size, train_size + test_size)) + print("init dataloader") + trainloader = torch.utils.data.DataLoader( + ds_train, + batch_size=args.mb, + shuffle=True, + num_workers=args.workers, + collate_fn=collate_fn_transformer_filter, + ) + testloader = torch.utils.data.DataLoader( + ds_test, + batch_size=args.mb, + shuffle=True, + num_workers=args.workers, + collate_fn=collate_fn_transformer_filter, + ) + print("init network") + models = [] + + if True: + for n_layers in args.n_layers: + models.append( + { + "name": "TrajectoryNet_l%d" % n_layers, + "model": TrajectoryNet( + d_drone_state=3 + 4, + d_radio_feature=args.beam_former_spacing + + 1, # add one for mean power normalization + d_detector_observation_embedding=args.obv_dembed, + d_trajectory_embedding=args.tj_dembed, + trajectory_prediction_n_layers=8, + d_trajectory_prediction_output=(2 + 2 + 1) + (2 + 2 + 1), + d_model=args.transformer_dmodel, + n_heads=8, + d_hid=64, + n_layers=n_layers, + n_outputs=8, + ssn_d_hid=args.ssn_dhid, + ssn_n_layers=args.ssn_nlayers, + # ssn_d_output=5 + ), + "fig": plt.figure(figsize=(18, 4)), + "dead": False, + "lr": args.lr_transformer, + } + ) + if False: + for n_layers in [2, 4, 8, 16, 32]: # ,32,64]: #,32,64]: + for snapshots_per_sample in args.snapshots_per_sample: + if snapshots_per_sample == 1: + continue + models.append( + { + "name": "%d snapshots (l%d)" % (snapshots_per_sample, n_layers), + "model": SnapshotNet( + snapshots_per_sample, + n_layers=n_layers, + d_model=args.transformer_dmodel, + n_outputs=len(cols_for_loss), + ssn_n_outputs=len(cols_for_loss), + dropout=0.0, + positional_encoding_len=args.positional_encoding_len, + tformer_input=args.transformer_input, + ), + "snapshots_per_sample": snapshots_per_sample, + "images": False, + "lr": args.lr_transformer, + "images": False, + "fig": plt.figure(figsize=(18, 4)), + "dead": False, + } + ) + + # move the models to the device + for d_model in models: + d_model["model"] = d_model["model"].to(dtype).to(device) + loss_figs = { + "train": plt.figure(figsize=(14 * 3, 6)), + "test": plt.figure(figsize=(14 * 3, 6)), + } + + for d_model in models: + d_model["optimizer"] = optim.Adam( + d_model["model"].parameters(), + lr=d_model["lr"], + weight_decay=args.weight_decay, + ) + if args.lr_scheduler_every != 0: + d_model["scheduler"] = optim.lr_scheduler.LinearLR( + d_model["optimizer"], + start_factor=0.001, + end_factor=1.0, + total_iters=30, + verbose=False, + ) - #move the models to the device - for d_model in models: - d_model['model']=d_model['model'].to(dtype).to(device) - loss_figs={ - 'train':plt.figure(figsize=(14*3,6)), - 'test':plt.figure(figsize=(14*3,6))} - - for d_model in models: - d_model['optimizer']=optim.Adam(d_model['model'].parameters(),lr=d_model['lr'],weight_decay=args.weight_decay) - if args.lr_scheduler_every!=0: - d_model['scheduler']=optim.lr_scheduler.LinearLR( - d_model['optimizer'], - start_factor=0.001, - end_factor=1.0, - total_iters=30, - verbose=False) - - criterion = nn.MSELoss().to(device) - - print("training loop") - running_losses={'train':{},'test':{}} - for k in ['train','test']: - running_losses[k]={ d['name']:[] for d in models} - running_losses[k]['baseline']=[] - - saves=[] - - start_time=time.time() - - def prep_data(data): - #todo add data augmentation - #split one source to two (merge) - #drop sources from input (birth) - #add sources not seen (death) - d={ k:data[k].to(dtype).to(device) for k in data} - - batch_size,time_steps,n_sources,_=d['emitter_position_and_velocity'].shape - d['emitter_position_and_velocity']=torch.cat([ - d['emitter_position_and_velocity'], - #torch.zeros(batch_size,time_steps,n_sources,4,device=d['emitter_position_and_velocity'].device)+1, # the variance - #torch.zeros(batch_size,time_steps,n_sources,2,device=d['emitter_position_and_velocity'].device), # the angle - ],dim=3) - for k in d: - assert(not d[k].isnan().any()) - #for k in d: - # print(k,(d[k]).mean(),torch.abs(np.abs(d[k])).mean()) - return d - - test_iterator = iter(testloader) - total_batch_idx=0 - for epoch in range(args.epochs): - for i, data in enumerate(trainloader, 0): - #move to device, do final prep - print(epoch,i) - total_batch_idx+=1 - prepared_data=prep_data(data) - if True: #torch.cuda.amp.autocast(): - for d_model in models: - d_model['model'].train() - for p in d_model['model'].parameters(): - if p.isnan().any(): - breakpoint() - loss,losses=model_forward( - d_model, - prepared_data, - args, - 'train', - update=total_batch_idx, - plot=True) - if args.lr_scheduler_every!=0 and total_batch_idx%args.lr_scheduler_every==args.lr_scheduler_every-1: - d_model['scheduler'].step() - loss.backward() - running_losses['train'][d_model['name']].append(losses) - if args.clip>0: - torch.nn.utils.clip_grad_norm_(d_model['model'].parameters(), args.clip) # clip gradients - d_model['optimizer'].step() - for p in d_model['model'].parameters(): - if p.isnan().any(): - breakpoint() - #labels=prepared_data['labels'] - running_losses['train']['baseline'].append({'baseline':1e-5}) - #running_losses['train']['baseline'].append({'baseline':criterion(labels*0+labels.mean(axis=[0,1],keepdim=True), labels).item() } ) - if total_batch_idx%args.print_every==args.print_every-1: - for idx in np.arange(args.test_mbs): - try: - data = next(test_iterator) - except StopIteration: - test_iterator = iter(testloader) - data = next(test_iterator) - prepared_data=prep_data(data) - with torch.no_grad(): - for d_model in models: - d_model['model'].eval() - loss,losses=model_forward( - d_model, - prepared_data, - args, - 'test', - update=total_batch_idx, - plot=idx==0) - running_losses['test'][d_model['name']].append(losses) - #labels=prepared_data['labels'] - running_losses['test']['baseline'].append({'baseline':1e-5})# {'baseline':criterion(labels*0+labels.mean(axis=[0,1],keepdim=True), labels).item() } ) - - - if total_batch_idx==0 or total_batch_idx%args.save_every==args.save_every-1: - save(args,running_losses,models,i,args.keep_n_saves) - - if total_batch_idx//args.print_every>=1 and (total_batch_idx % args.print_every) == args.print_every-1: - train_baseline_loss=model_to_losses(running_losses['train']['baseline'],args.print_every) - test_baseline_loss=model_to_losses(running_losses['test']['baseline'],args.test_mbs) - print(f'[{epoch + 1}, {i + 1:5d}]') - print(f'\tTrain: baseline: {train_baseline_loss["baseline"][-1]:.3f} , time { (time.time()-start_time)/(i+1) :.3f} / batch' ) - print(f'\tTest: baseline: {test_baseline_loss["baseline"][-1]:.3f}, time { (time.time()-start_time)/(i+1) :.3f} / batch' ) - loss_str="\t"+"\n\t".join( - [ "%s(%s):(tr)%s,(ts)%s" % (d['name'],str(d['dead']), - model_to_loss_str(running_losses['train'][d['name']],args.print_every), - model_to_loss_str(running_losses['test'][d['name']],args.test_mbs) - ) for d in models ]) - print(loss_str) - - if total_batch_idx//args.print_every>=1 and (total_batch_idx % args.plot_every) == args.plot_every-1: - plot_loss(running_losses=running_losses['train'], - baseline_loss=train_baseline_loss, - xtick_spacing=args.print_every, - mean_chunk=args.print_every, - output_prefix=args.output_prefix, - fig=loss_figs['train'], - title='Train',update=i) - plot_loss(running_losses=running_losses['test'], - baseline_loss=test_baseline_loss, - xtick_spacing=args.print_every, - mean_chunk=args.test_mbs, - output_prefix=args.output_prefix, - fig=loss_figs['test'], - title='Test',update=i) - if args.plot: - plt.pause(0.5) - - print('Finished Training') # but do we ever really get here? + criterion = nn.MSELoss().to(device) + + print("training loop") + running_losses = {"train": {}, "test": {}} + for k in ["train", "test"]: + running_losses[k] = {d["name"]: [] for d in models} + running_losses[k]["baseline"] = [] + + saves = [] + + start_time = time.time() + + def prep_data(data): + # todo add data augmentation + # split one source to two (merge) + # drop sources from input (birth) + # add sources not seen (death) + d = {k: data[k].to(dtype).to(device) for k in data} + + batch_size, time_steps, n_sources, _ = d["emitter_position_and_velocity"].shape + d["emitter_position_and_velocity"] = torch.cat( + [ + d["emitter_position_and_velocity"], + # torch.zeros(batch_size,time_steps,n_sources,4,device=d['emitter_position_and_velocity'].device)+1, # the variance + # torch.zeros(batch_size,time_steps,n_sources,2,device=d['emitter_position_and_velocity'].device), # the angle + ], + dim=3, + ) + for k in d: + assert not d[k].isnan().any() + # for k in d: + # print(k,(d[k]).mean(),torch.abs(np.abs(d[k])).mean()) + return d + + test_iterator = iter(testloader) + total_batch_idx = 0 + for epoch in range(args.epochs): + for i, data in enumerate(trainloader, 0): + # move to device, do final prep + print(epoch, i) + total_batch_idx += 1 + prepared_data = prep_data(data) + if True: # torch.cuda.amp.autocast(): + for d_model in models: + d_model["model"].train() + for p in d_model["model"].parameters(): + if p.isnan().any(): + breakpoint() + loss, losses = model_forward( + d_model, + prepared_data, + args, + "train", + update=total_batch_idx, + plot=True, + ) + if ( + args.lr_scheduler_every != 0 + and total_batch_idx % args.lr_scheduler_every + == args.lr_scheduler_every - 1 + ): + d_model["scheduler"].step() + loss.backward() + running_losses["train"][d_model["name"]].append(losses) + if args.clip > 0: + torch.nn.utils.clip_grad_norm_( + d_model["model"].parameters(), args.clip + ) # clip gradients + d_model["optimizer"].step() + for p in d_model["model"].parameters(): + if p.isnan().any(): + breakpoint() + # labels=prepared_data['labels'] + running_losses["train"]["baseline"].append({"baseline": 1e-5}) + # running_losses['train']['baseline'].append({'baseline':criterion(labels*0+labels.mean(axis=[0,1],keepdim=True), labels).item() } ) + if total_batch_idx % args.print_every == args.print_every - 1: + for idx in np.arange(args.test_mbs): + try: + data = next(test_iterator) + except StopIteration: + test_iterator = iter(testloader) + data = next(test_iterator) + prepared_data = prep_data(data) + with torch.no_grad(): + for d_model in models: + d_model["model"].eval() + loss, losses = model_forward( + d_model, + prepared_data, + args, + "test", + update=total_batch_idx, + plot=idx == 0, + ) + running_losses["test"][d_model["name"]].append(losses) + # labels=prepared_data['labels'] + running_losses["test"]["baseline"].append( + {"baseline": 1e-5} + ) # {'baseline':criterion(labels*0+labels.mean(axis=[0,1],keepdim=True), labels).item() } ) + + if ( + total_batch_idx == 0 + or total_batch_idx % args.save_every == args.save_every - 1 + ): + save(args, running_losses, models, i, args.keep_n_saves) + + if ( + total_batch_idx // args.print_every >= 1 + and (total_batch_idx % args.print_every) == args.print_every - 1 + ): + train_baseline_loss = model_to_losses( + running_losses["train"]["baseline"], args.print_every + ) + test_baseline_loss = model_to_losses( + running_losses["test"]["baseline"], args.test_mbs + ) + print(f"[{epoch + 1}, {i + 1:5d}]") + print( + f'\tTrain: baseline: {train_baseline_loss["baseline"][-1]:.3f} , time { (time.time()-start_time)/(i+1) :.3f} / batch' + ) + print( + f'\tTest: baseline: {test_baseline_loss["baseline"][-1]:.3f}, time { (time.time()-start_time)/(i+1) :.3f} / batch' + ) + loss_str = "\t" + "\n\t".join( + [ + "%s(%s):(tr)%s,(ts)%s" + % ( + d["name"], + str(d["dead"]), + model_to_loss_str( + running_losses["train"][d["name"]], args.print_every + ), + model_to_loss_str( + running_losses["test"][d["name"]], args.test_mbs + ), + ) + for d in models + ] + ) + print(loss_str) + + if ( + total_batch_idx // args.print_every >= 1 + and (total_batch_idx % args.plot_every) == args.plot_every - 1 + ): + plot_loss( + running_losses=running_losses["train"], + baseline_loss=train_baseline_loss, + xtick_spacing=args.print_every, + mean_chunk=args.print_every, + output_prefix=args.output_prefix, + fig=loss_figs["train"], + title="Train", + update=i, + ) + plot_loss( + running_losses=running_losses["test"], + baseline_loss=test_baseline_loss, + xtick_spacing=args.print_every, + mean_chunk=args.test_mbs, + output_prefix=args.output_prefix, + fig=loss_figs["test"], + title="Test", + update=i, + ) + if args.plot: + plt.pause(0.5) + + print("Finished Training") # but do we ever really get here? diff --git a/software/model_training_and_inference/90_real_session_plotter.py b/software/model_training_and_inference/90_real_session_plotter.py index 2f76ff03..125fbdc4 100644 --- a/software/model_training_and_inference/90_real_session_plotter.py +++ b/software/model_training_and_inference/90_real_session_plotter.py @@ -5,35 +5,35 @@ from compress_pickle import dump, load from utils.spf_dataset import SessionsDatasetReal -if __name__=='__main__': - parser = argparse.ArgumentParser() - parser.add_argument('--root-dir', type=str, required=True) - parser.add_argument('--snapshots-in-file', type=int, required=False,default=400000) - parser.add_argument('--nthetas', type=int, required=False,default=64+1) - parser.add_argument('--snapshots-in-sample', type=int, required=False,default=128) - parser.add_argument('--width', type=int, required=False,default=3000) - parser.add_argument('--session-idx', type=int, required=True) - parser.add_argument('--steps', type=int, required=False, default=3) - parser.add_argument('--step-size', type=int, required=False, default=16) - parser.add_argument('--duration', type=int, required=False, default=50) - parser.add_argument('--output_prefix', type=str, required=False, default='session_output') - args = parser.parse_args() - assert(args.snapshots_in_sample>=args.steps) - ds=SessionsDatasetReal( - root_dir=args.root_dir, - snapshots_in_file=args.snapshots_in_file, - nthetas=args.nthetas, - snapshots_in_sample=args.snapshots_in_sample, - nsources=1, - width=args.width, - step_size=args.step_size - ) - print(ds) - print("DSLEN",ds.len) - - session=ds[args.session_idx] - filenames=plot_full_session(session,args.steps,args.output_prefix,invert=True) - - filenames_to_gif(filenames,"%s.gif" % args.output_prefix,duration=args.duration) +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--root-dir", type=str, required=True) + parser.add_argument("--snapshots-in-file", type=int, required=False, default=400000) + parser.add_argument("--nthetas", type=int, required=False, default=64 + 1) + parser.add_argument("--snapshots-in-sample", type=int, required=False, default=128) + parser.add_argument("--width", type=int, required=False, default=3000) + parser.add_argument("--session-idx", type=int, required=True) + parser.add_argument("--steps", type=int, required=False, default=3) + parser.add_argument("--step-size", type=int, required=False, default=16) + parser.add_argument("--duration", type=int, required=False, default=50) + parser.add_argument( + "--output_prefix", type=str, required=False, default="session_output" + ) + args = parser.parse_args() + assert args.snapshots_in_sample >= args.steps + ds = SessionsDatasetReal( + root_dir=args.root_dir, + snapshots_in_file=args.snapshots_in_file, + nthetas=args.nthetas, + snapshots_in_sample=args.snapshots_in_sample, + nsources=1, + width=args.width, + step_size=args.step_size, + ) + print(ds) + print("DSLEN", ds.len) + session = ds[args.session_idx] + filenames = plot_full_session(session, args.steps, args.output_prefix, invert=True) + filenames_to_gif(filenames, "%s.gif" % args.output_prefix, duration=args.duration) diff --git a/software/model_training_and_inference/90_session_plotter.py b/software/model_training_and_inference/90_session_plotter.py index 27094cec..68cfa8d3 100644 --- a/software/model_training_and_inference/90_session_plotter.py +++ b/software/model_training_and_inference/90_session_plotter.py @@ -5,18 +5,18 @@ from compress_pickle import dump, load from utils.spf_dataset import SessionsDatasetTask2 -if __name__=='__main__': - parser = argparse.ArgumentParser() - parser.add_argument('--dataset', type=str, required=True) - parser.add_argument('--session-idx', type=int, required=True) - parser.add_argument('--steps', type=int, required=False, default=3) - parser.add_argument('--output_prefix', type=str, required=False, default='session_output') - args = parser.parse_args() - - ds=SessionsDatasetTask2(args.dataset,snapshots_in_sample=args.steps) - session=ds[args.session_idx] - filenames=plot_full_session(session,args.steps,args.output_prefix) - - filenames_to_gif(filenames,"%s.gif" % args.output_prefix) +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--dataset", type=str, required=True) + parser.add_argument("--session-idx", type=int, required=True) + parser.add_argument("--steps", type=int, required=False, default=3) + parser.add_argument( + "--output_prefix", type=str, required=False, default="session_output" + ) + args = parser.parse_args() + ds = SessionsDatasetTask2(args.dataset, snapshots_in_sample=args.steps) + session = ds[args.session_idx] + filenames = plot_full_session(session, args.steps, args.output_prefix) + filenames_to_gif(filenames, "%s.gif" % args.output_prefix) diff --git a/software/model_training_and_inference/91_line_plotter.py b/software/model_training_and_inference/91_line_plotter.py index dfc9ee8a..2e419673 100644 --- a/software/model_training_and_inference/91_line_plotter.py +++ b/software/model_training_and_inference/91_line_plotter.py @@ -5,18 +5,18 @@ from compress_pickle import dump, load from utils.spf_dataset import SessionsDataset -if __name__=='__main__': - parser = argparse.ArgumentParser() - parser.add_argument('--dataset', type=str, required=True) - parser.add_argument('--session-idx', type=int, required=True) - parser.add_argument('--steps', type=int, required=False, default=3) - parser.add_argument('--output_prefix', type=str, required=False, default='session_output') - args = parser.parse_args() - - ds=SessionsDataset(args.dataset,snapshots_in_sample=args.steps) - session=ds[args.session_idx] - filenames=plot_lines(session,args.steps,args.output_prefix) - - filenames_to_gif(filenames,"%s_lines.gif" % args.output_prefix,size=(1200,400)) +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--dataset", type=str, required=True) + parser.add_argument("--session-idx", type=int, required=True) + parser.add_argument("--steps", type=int, required=False, default=3) + parser.add_argument( + "--output_prefix", type=str, required=False, default="session_output" + ) + args = parser.parse_args() + ds = SessionsDataset(args.dataset, snapshots_in_sample=args.steps) + session = ds[args.session_idx] + filenames = plot_lines(session, args.steps, args.output_prefix) + filenames_to_gif(filenames, "%s_lines.gif" % args.output_prefix, size=(1200, 400)) diff --git a/software/model_training_and_inference/92_evaluate_session.py b/software/model_training_and_inference/92_evaluate_session.py index 2b85cc61..d96512e5 100644 --- a/software/model_training_and_inference/92_evaluate_session.py +++ b/software/model_training_and_inference/92_evaluate_session.py @@ -3,66 +3,69 @@ import numpy as np from utils.plot import filenames_to_gif, plot_lines, plot_predictions_and_baseline from compress_pickle import dump, load -from utils.spf_dataset import SessionsDatasetTask2,collate_fn +from utils.spf_dataset import SessionsDatasetTask2, collate_fn import torch import random -from utils.baseline_algorithm import baseline_algorithm +from utils.baseline_algorithm import baseline_algorithm from utils.spf_dataset import collate_fn, output_cols, rel_to_pos from utils.save_load import CPU_Unpickler -if __name__=='__main__': - parser = argparse.ArgumentParser() - parser.add_argument('--dataset', type=str, required=True) - parser.add_argument('--session-idx', type=int, required=True) - parser.add_argument('--seed', type=int, required=False, default=0) - parser.add_argument('--load', type=str, required=True) - parser.add_argument('--test-fraction', type=float, required=False, default=0.2) - parser.add_argument('--model-name', type=str, required=True) +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--dataset", type=str, required=True) + parser.add_argument("--session-idx", type=int, required=True) + parser.add_argument("--seed", type=int, required=False, default=0) + parser.add_argument("--load", type=str, required=True) + parser.add_argument("--test-fraction", type=float, required=False, default=0.2) + parser.add_argument("--model-name", type=str, required=True) - parser.add_argument('--output_prefix', type=str, required=False, default='session_output') - args = parser.parse_args() + parser.add_argument( + "--output_prefix", type=str, required=False, default="session_output" + ) + args = parser.parse_args() - d=CPU_Unpickler(open(args.load,'rb')).load() - str_to_model={ model['name']:model for model in d['models'] } - model=str_to_model[args.model_name] + d = CPU_Unpickler(open(args.load, "rb")).load() + str_to_model = {model["name"]: model for model in d["models"]} + model = str_to_model[args.model_name] - args.snapshots_per_sample=model['snapshots_per_sample'] + args.snapshots_per_sample = model["snapshots_per_sample"] - torch.manual_seed(args.seed) - random.seed(args.seed) - np.random.seed(args.seed) - - ds=SessionsDatasetTask2(args.dataset,snapshots_in_sample=args.snapshots_per_sample) - train_size=int(len(ds)*args.test_fraction) - test_size=len(ds)-train_size + torch.manual_seed(args.seed) + random.seed(args.seed) + np.random.seed(args.seed) - ds_train = torch.utils.data.Subset(ds, np.arange(train_size)) - ds_test = torch.utils.data.Subset(ds, np.arange(train_size, train_size + test_size)) + ds = SessionsDatasetTask2( + args.dataset, snapshots_in_sample=args.snapshots_per_sample + ) + train_size = int(len(ds) * args.test_fraction) + test_size = len(ds) - train_size - session=ds_test[args.session_idx] - #_in=collate_fn([session]) - _in=collate_fn([session]) - width=session['width_at_t'][0][0] + ds_train = torch.utils.data.Subset(ds, np.arange(train_size)) + ds_test = torch.utils.data.Subset(ds, np.arange(train_size, train_size + test_size)) - # get the baseline - fp,imgs,baseline_predictions=baseline_algorithm(session,width,steps=model['snapshots_per_sample']) - - # get the predictions from the model - with torch.no_grad(): - model_pred=model['model'].cpu()(_in['inputs'])['transformer_pred'] - model_predictions=rel_to_pos(model_pred[0,:,output_cols['src_pos']],width).clamp(0,width) - plot_predictions_and_baseline(session,args,model['snapshots_per_sample']-1, - { - 'name':'baseline algorithm', - 'predictions':baseline_predictions - }, - { - 'name':'NN '+args.model_name, - 'predictions':model_predictions - }, - ) - #filenames=plot_lines(session,args.steps,args.output_prefix) + session = ds_test[args.session_idx] + # _in=collate_fn([session]) + _in = collate_fn([session]) + width = session["width_at_t"][0][0] - #filenames_to_gif(filenames,"%s_lines.gif" % args.output_prefix,size=(1200,400)) + # get the baseline + fp, imgs, baseline_predictions = baseline_algorithm( + session, width, steps=model["snapshots_per_sample"] + ) + # get the predictions from the model + with torch.no_grad(): + model_pred = model["model"].cpu()(_in["inputs"])["transformer_pred"] + model_predictions = rel_to_pos( + model_pred[0, :, output_cols["src_pos"]], width + ).clamp(0, width) + plot_predictions_and_baseline( + session, + args, + model["snapshots_per_sample"] - 1, + {"name": "baseline algorithm", "predictions": baseline_predictions}, + {"name": "NN " + args.model_name, "predictions": model_predictions}, + ) + # filenames=plot_lines(session,args.steps,args.output_prefix) + # filenames_to_gif(filenames,"%s_lines.gif" % args.output_prefix,size=(1200,400)) diff --git a/software/model_training_and_inference/models/models.py b/software/model_training_and_inference/models/models.py index 30623651..0be03294 100644 --- a/software/model_training_and_inference/models/models.py +++ b/software/model_training_and_inference/models/models.py @@ -13,711 +13,816 @@ from torch.utils.data import dataset import math + class PositionalEncoding(nn.Module): + def __init__(self, d, l): + super().__init__() - def __init__(self, d, l): - super().__init__() + position = torch.arange(l).unsqueeze(1) + div_term = torch.exp(torch.arange(0, d, 2) * (-math.log(10000.0) / d)) + pe = torch.zeros(1, l, d) + pe[0, :, 0::2] = torch.sin(position * div_term) + pe[0, :, 1::2] = torch.cos(position * div_term) + self.register_buffer("pe", pe) - position = torch.arange(l).unsqueeze(1) - div_term = torch.exp(torch.arange(0, d, 2) * (-math.log(10000.0) / d)) - pe = torch.zeros(1,l, d) - pe[0, :, 0::2] = torch.sin(position * div_term) - pe[0, :, 1::2] = torch.cos(position * div_term) - self.register_buffer('pe', pe) + def forward(self, x): + """ + Arguments: + x: Tensor, shape ``[batch, time_steps, features]`` + """ + _pe = self.pe.expand((x.shape[0], self.pe.shape[1], self.pe.shape[2])) + return torch.cat([x, _pe], axis=2) - def forward(self, x): - """ - Arguments: - x: Tensor, shape ``[batch, time_steps, features]`` - """ - _pe=self.pe.expand((x.shape[0],self.pe.shape[1],self.pe.shape[2])) - return torch.cat([x,_pe],axis=2) if False: - from complexPyTorch.complexLayers import ComplexBatchNorm2d, ComplexConv2d, ComplexLinear, ComplexReLU, ComplexBatchNorm1d #, ComplexSigmoid - from complexPyTorch.complexFunctions import complex_relu - class HybridFFNN(nn.Module): - def __init__(self, - d_inputs, - d_outputs, - n_complex_layers, - n_real_layers, - d_hidden, - norm=False): - super().__init__() - - self.complex_net=ComplexFFNN( - d_inputs, - d_hidden, - n_layers=n_complex_layers, - d_hidden=d_hidden, - norm=norm - ) - self.real_net=nn.Sequential( - nn.Linear(d_hidden,d_hidden), - *[nn.Sequential( - nn.Linear(d_hidden,d_hidden), - nn.LayerNorm(d_hidden),# if norm else nn.Identity(), - nn.ReLU() - ) - for _ in range(n_real_layers) ], - nn.LayerNorm(d_hidden),# if norm else nn.Identity(), - nn.Linear(d_hidden,d_outputs) - ) - - def forward(self,x): - complex_out=self.complex_net(x) - real_out=self.real_net(complex_out.abs()) - return F.softmax(real_out,dim=1) - #return real_out/(real_out.sum(axis=1,keepdims=True)+1e-9) - - - class ComplexFFNN(nn.Module): - def __init__(self, - d_inputs, - d_outputs, - n_layers, - d_hidden, - norm=False): - super().__init__() - - self.output_net=nn.Sequential( - ComplexBatchNorm1d(d_inputs) if norm else nn.Identity(), - ComplexLinear(d_inputs,d_hidden), - *[nn.Sequential( - ComplexLinear(d_hidden,d_hidden), - ComplexReLU(), - ComplexBatchNorm1d(d_hidden) if norm else nn.Identity(), - ) - for _ in range(n_layers) ], - ComplexLinear(d_hidden,d_outputs), - ) - def forward(self,x): - out=self.output_net(x).abs() - if out.sum().isnan(): - breakpoint() - a=1 - #breakpoint() - return F.softmax(self.output_net(x).abs(),dim=1) - #return out/(out.sum(axis=1,keepdims=True)+1e-9) - - + from complexPyTorch.complexLayers import ( + ComplexBatchNorm2d, + ComplexConv2d, + ComplexLinear, + ComplexReLU, + ComplexBatchNorm1d, + ) # , ComplexSigmoid + from complexPyTorch.complexFunctions import complex_relu + + class HybridFFNN(nn.Module): + def __init__( + self, + d_inputs, + d_outputs, + n_complex_layers, + n_real_layers, + d_hidden, + norm=False, + ): + super().__init__() + + self.complex_net = ComplexFFNN( + d_inputs, + d_hidden, + n_layers=n_complex_layers, + d_hidden=d_hidden, + norm=norm, + ) + self.real_net = nn.Sequential( + nn.Linear(d_hidden, d_hidden), + *[ + nn.Sequential( + nn.Linear(d_hidden, d_hidden), + nn.LayerNorm(d_hidden), # if norm else nn.Identity(), + nn.ReLU(), + ) + for _ in range(n_real_layers) + ], + nn.LayerNorm(d_hidden), # if norm else nn.Identity(), + nn.Linear(d_hidden, d_outputs), + ) + + def forward(self, x): + complex_out = self.complex_net(x) + real_out = self.real_net(complex_out.abs()) + return F.softmax(real_out, dim=1) + # return real_out/(real_out.sum(axis=1,keepdims=True)+1e-9) + + class ComplexFFNN(nn.Module): + def __init__(self, d_inputs, d_outputs, n_layers, d_hidden, norm=False): + super().__init__() + + self.output_net = nn.Sequential( + ComplexBatchNorm1d(d_inputs) if norm else nn.Identity(), + ComplexLinear(d_inputs, d_hidden), + *[ + nn.Sequential( + ComplexLinear(d_hidden, d_hidden), + ComplexReLU(), + ComplexBatchNorm1d(d_hidden) if norm else nn.Identity(), + ) + for _ in range(n_layers) + ], + ComplexLinear(d_hidden, d_outputs), + ) + + def forward(self, x): + out = self.output_net(x).abs() + if out.sum().isnan(): + breakpoint() + a = 1 + # breakpoint() + return F.softmax(self.output_net(x).abs(), dim=1) + # return out/(out.sum(axis=1,keepdims=True)+1e-9) + class TransformerEncOnlyModel(nn.Module): - def __init__(self, - d_in, - d_model, - n_heads, - d_hid, - n_layers, - dropout, - n_outputs, - n_layers_output=4): - super().__init__() - self.model_type = 'Transformer' - - encoder_layers = TransformerEncoderLayer( - d_model, - n_heads, - d_hid, - dropout, - activation='gelu', - batch_first=True) - self.transformer_encoder = TransformerEncoder( - encoder_layers, - n_layers, - nn.LayerNorm(d_model), - ) - - assert( d_model>=d_in) - - self.linear_in = nn.Linear(d_in, d_model) if d_model>d_in else nn.Identity() - self.d_model=d_model - - self.output_net=nn.Sequential( - nn.Linear(self.d_model,d_hid), - nn.SELU(), - *[SkipConnection(nn.Sequential( - nn.Linear(d_hid,d_hid), - nn.LayerNorm(d_hid), - nn.SELU() - )) - for _ in range(n_layers_output) ], - #nn.LayerNorm(d_hid), - nn.Linear(d_hid,n_outputs), - nn.LayerNorm(n_outputs) - ) - - def forward(self, src: Tensor, src_key_padding_mask=None) -> Tensor: - #src_enc = self.transformer_encoder( - # torch.cat( - # [src,self.linear_in(src)],axis=2)) #/np.sqrt(self.d_radio_feature)) - src_enc = self.transformer_encoder( self.linear_in(src), src_key_padding_mask=src_key_padding_mask) - #output = self.transformer_encoder(src) #,self.linear_in(src)],axis=2)) #/np.sqrt(self.d_radio_feature)) - output = self.output_net(src_enc) #/np.sqrt(self.d_model) - if output.isnan().any(): - breakpoint() - return output + def __init__( + self, + d_in, + d_model, + n_heads, + d_hid, + n_layers, + dropout, + n_outputs, + n_layers_output=4, + ): + super().__init__() + self.model_type = "Transformer" + + encoder_layers = TransformerEncoderLayer( + d_model, n_heads, d_hid, dropout, activation="gelu", batch_first=True + ) + self.transformer_encoder = TransformerEncoder( + encoder_layers, + n_layers, + nn.LayerNorm(d_model), + ) + + assert d_model >= d_in + + self.linear_in = nn.Linear(d_in, d_model) if d_model > d_in else nn.Identity() + self.d_model = d_model + + self.output_net = nn.Sequential( + nn.Linear(self.d_model, d_hid), + nn.SELU(), + *[ + SkipConnection( + nn.Sequential( + nn.Linear(d_hid, d_hid), nn.LayerNorm(d_hid), nn.SELU() + ) + ) + for _ in range(n_layers_output) + ], + # nn.LayerNorm(d_hid), + nn.Linear(d_hid, n_outputs), + nn.LayerNorm(n_outputs), + ) + + def forward(self, src: Tensor, src_key_padding_mask=None) -> Tensor: + # src_enc = self.transformer_encoder( + # torch.cat( + # [src,self.linear_in(src)],axis=2)) #/np.sqrt(self.d_radio_feature)) + src_enc = self.transformer_encoder( + self.linear_in(src), src_key_padding_mask=src_key_padding_mask + ) + # output = self.transformer_encoder(src) #,self.linear_in(src)],axis=2)) #/np.sqrt(self.d_radio_feature)) + output = self.output_net(src_enc) # /np.sqrt(self.d_model) + if output.isnan().any(): + breakpoint() + return output class SkipConnection(nn.Module): - def __init__(self,module): - super().__init__() - self.module=module + def __init__(self, module): + super().__init__() + self.module = module + + def forward(self, x): + return self.module(x) + x - def forward(self,x): - return self.module(x)+x class FilterNet(nn.Module): - def __init__(self, - d_drone_state=4+4, - d_emitter_state=4+4+2, # 2 means , 2 variances, 2 angles - d_radio_feature=258, - d_model=512, - n_heads=8, - d_hid=256, - d_embed=64, - n_layers=1, - n_outputs=4+4+2+1, # 2 means , 2 variances, 2 angles, responsibility - dropout=0.0, - ssn_d_hid=64, - ssn_n_layers=8, - ssn_n_outputs=8, - ssn_dropout=0.0, - tformer_input=['drone_state','embedding']): - super().__init__() - self.d_radio_feature=d_radio_feature - self.d_drone_state=d_drone_state - self.d_model=d_model - self.n_heads=n_heads - self.d_embed=d_embed - self.d_hid=d_hid - self.dropout=dropout - self.tformer_input=tformer_input - self.tformer_input_dim=0 - self.n_layers=n_layers - self.n_outputs=n_outputs - - - for k,v in [ - ('drone_state',d_drone_state), - ('radio_feature',d_radio_feature), - ('embedding',self.d_embed)]: - if k in self.tformer_input: - self.tformer_input_dim+=v - - self.snap_shot_net=SingleSnapshotNet( - d_input_feature=d_radio_feature+d_drone_state, - d_hid=ssn_d_hid, - d_embed=self.d_embed-1, - n_layers=ssn_n_layers, - n_outputs=n_outputs, - dropout=ssn_dropout) - - self.transformer_enc=TransformerEncOnlyModel( - d_in=self.d_model, - d_model=self.d_model, - n_heads=self.n_heads, - d_hid=self.d_hid, - n_layers=self.n_layers, - dropout=self.dropout, - n_outputs=self.n_outputs, - n_layers_output=4) - - def forward(self,x): - single_snapshot_output=self.snap_shot_net(x) - snap_shot_embeddings=single_snapshot_output['embedding'] #torch.Size([batch, time, embedding dim]) - emitter_state_embeddings=self.emitter_embedding_net(x)['emitter_state_embedding'] - - batch_size,time_steps,n_sources,_=emitter_state_embeddings.shape - - snap_shot_embeddings=torch.cat([ - snap_shot_embeddings, #torch.Size([batch, time, embedding dim]) - torch.zeros(batch_size,time_steps,1).to(snap_shot_embeddings.device)],dim=2).reshape(batch_size,time_steps,1,self.d_embed) - emitter_state_embeddings=torch.cat([ - emitter_state_embeddings, - torch.ones(batch_size,time_steps,n_sources,1).to(emitter_state_embeddings.device)],dim=3) - - - self_attention_input=torch.cat([emitter_state_embeddings,snap_shot_embeddings],dim=2) - #input shape = (batch*time, nsources+1snapshot, embedding_dim) - #output shape = (batch*time,N,embedding_dim) - - self_attention_output=self.transformer_enc( - self_attention_input.reshape( - batch_size*time_steps, - n_sources+1, - self.d_model)).reshape(batch_size,time_steps,n_sources+1,self.n_outputs) - #make sure assignments are probabilities? - self_attention_output=torch.cat([self_attention_output[...,:-1],nn.functional.softmax(self_attention_output[...,[-1]],dim=2)],dim=3) - return {'transformer_pred':self_attention_output, - 'single_snapshot_pred':single_snapshot_output['single_snapshot_pred']} + def __init__( + self, + d_drone_state=4 + 4, + d_emitter_state=4 + 4 + 2, # 2 means , 2 variances, 2 angles + d_radio_feature=258, + d_model=512, + n_heads=8, + d_hid=256, + d_embed=64, + n_layers=1, + n_outputs=4 + 4 + 2 + 1, # 2 means , 2 variances, 2 angles, responsibility + dropout=0.0, + ssn_d_hid=64, + ssn_n_layers=8, + ssn_n_outputs=8, + ssn_dropout=0.0, + tformer_input=["drone_state", "embedding"], + ): + super().__init__() + self.d_radio_feature = d_radio_feature + self.d_drone_state = d_drone_state + self.d_model = d_model + self.n_heads = n_heads + self.d_embed = d_embed + self.d_hid = d_hid + self.dropout = dropout + self.tformer_input = tformer_input + self.tformer_input_dim = 0 + self.n_layers = n_layers + self.n_outputs = n_outputs + + for k, v in [ + ("drone_state", d_drone_state), + ("radio_feature", d_radio_feature), + ("embedding", self.d_embed), + ]: + if k in self.tformer_input: + self.tformer_input_dim += v + + self.snap_shot_net = SingleSnapshotNet( + d_input_feature=d_radio_feature + d_drone_state, + d_hid=ssn_d_hid, + d_embed=self.d_embed - 1, + n_layers=ssn_n_layers, + n_outputs=n_outputs, + dropout=ssn_dropout, + ) + + self.transformer_enc = TransformerEncOnlyModel( + d_in=self.d_model, + d_model=self.d_model, + n_heads=self.n_heads, + d_hid=self.d_hid, + n_layers=self.n_layers, + dropout=self.dropout, + n_outputs=self.n_outputs, + n_layers_output=4, + ) + + def forward(self, x): + single_snapshot_output = self.snap_shot_net(x) + snap_shot_embeddings = single_snapshot_output[ + "embedding" + ] # torch.Size([batch, time, embedding dim]) + emitter_state_embeddings = self.emitter_embedding_net(x)[ + "emitter_state_embedding" + ] + + batch_size, time_steps, n_sources, _ = emitter_state_embeddings.shape + + snap_shot_embeddings = torch.cat( + [ + snap_shot_embeddings, # torch.Size([batch, time, embedding dim]) + torch.zeros(batch_size, time_steps, 1).to(snap_shot_embeddings.device), + ], + dim=2, + ).reshape(batch_size, time_steps, 1, self.d_embed) + emitter_state_embeddings = torch.cat( + [ + emitter_state_embeddings, + torch.ones(batch_size, time_steps, n_sources, 1).to( + emitter_state_embeddings.device + ), + ], + dim=3, + ) + + self_attention_input = torch.cat( + [emitter_state_embeddings, snap_shot_embeddings], dim=2 + ) + # input shape = (batch*time, nsources+1snapshot, embedding_dim) + # output shape = (batch*time,N,embedding_dim) + + self_attention_output = self.transformer_enc( + self_attention_input.reshape( + batch_size * time_steps, n_sources + 1, self.d_model + ) + ).reshape(batch_size, time_steps, n_sources + 1, self.n_outputs) + # make sure assignments are probabilities? + self_attention_output = torch.cat( + [ + self_attention_output[..., :-1], + nn.functional.softmax(self_attention_output[..., [-1]], dim=2), + ], + dim=3, + ) + return { + "transformer_pred": self_attention_output, + "single_snapshot_pred": single_snapshot_output["single_snapshot_pred"], + } + class TrajectoryNet(nn.Module): - def __init__(self, - d_drone_state=3+4, - d_radio_feature=258, - d_detector_observation_embedding=128, - d_trajectory_embedding=256, - trajectory_prediction_n_layers=8, - d_trajectory_prediction_output=(2+2+1)+(2+2+1), - d_model=512, - n_heads=8, - d_hid=256, - n_layers=4, - n_outputs=8, - ssn_d_hid=64, - ssn_n_layers=8, - #ssn_d_output=5, - ): - super().__init__() - self.d_detector_observation_embedding=d_detector_observation_embedding - self.d_trajectory_embedding=d_trajectory_embedding - self.d_trajectory_prediction_output=d_trajectory_prediction_output - - self.snap_shot_net=SnapShotEmbeddingNet( - d_in=d_radio_feature+d_drone_state, # time is inside drone state - d_hid=ssn_d_hid, - #d_out=5,# 2 mean, 2 sigma, 1 angle - d_embed=d_detector_observation_embedding, - n_layers=ssn_n_layers - ) - - self.trajectory_prediction_net=EmbeddingNet( - d_in=d_trajectory_embedding+1, - d_hid=d_trajectory_embedding, - d_out=d_trajectory_prediction_output, - d_embed=d_trajectory_embedding, - n_layers=trajectory_prediction_n_layers - ) - - self.tformer=TransformerEncOnlyModel( - d_in=d_detector_observation_embedding+1, - d_model=d_model, - n_heads=n_heads, - d_hid=d_hid, - n_layers=n_layers, - dropout=0.1, - n_outputs=d_trajectory_embedding) - - def l2(self): - return { - 'snap_shot_net': sum([ (p**2).sum() for p in self.snap_shot_net.parameters() ]) - } - - def forward(self,x): - # Lets get the detector observation embeddings - batch_size,time_steps,n_emitters,_=x['emitters_n_broadcasts'].shape - device=x['emitters_n_broadcasts'].device - - ############### - # SINGLE SNAPSHOT PREDICTIONS + EMBEDDINGS - ############### - d=self.snap_shot_net( - torch.cat([ # drone state + radio feature as input to SSN - x['drone_state'], - x['radio_feature'] - ],dim=2) - ) - drone_state_and_observation_embeddings=d['embedding'].detach() - single_snapshot_predictions=d['output'] - - - ################ - # TRAJECTORY EMBEDDINGS - #pick random time for each batch item, crop session, and compute tracjectory emebedding for each tracked object - ################ - #generate random times to grab - if self.training: - rt=torch.randint(low=max(2,batch_size//10), high=time_steps-1, size=(batch_size,)) # keep on CPU?, low=2 here gaurantees at least one thing was being tracked for each example - rt[:batch_size//4]=(time_steps-1)//2+1 # this might make it a bit smoother? - rt[batch_size//2:((3*batch_size)//4)]=time_steps # this might make it a bit smoother? - else: - print("EVAL MODE") - rt=torch.ones(batch_size,dtype=int)*(time_steps-1) - #now lets grab the (nsources,1) vector for each batch example that tells us how many times the object has transmitted previously - tracking=x['emitters_n_broadcasts'][torch.arange(batch_size),rt-1].cpu() #positive values for things already being tracked - - #find out how many objects are being tracked in each batch example - # so that we can pull out their embeddings to calculate trajectory embeddings - nested_tensors=[] - idxs=[] # List of batch idx and idx of object being tracked - max_snapshots=0 - for b in torch.arange(batch_size): - t=rt[b] - for tracked in torch.where(tracking[b]>0)[0]: - # get the mask where this emitter is broadcasting in the first t steps - times_where_this_tracked_is_broadcasting=x['emitters_broadcasting'][b,:t,tracked,0].to(bool) - # pull the drone state and observations for these time steps - #nested_tensors.append( - # drone_state_and_observation_embeddings[b,:t][times_where_this_tracked_is_broadcasting] - #) - nested_tensors.append( - torch.hstack([ - drone_state_and_observation_embeddings[b,:t][times_where_this_tracked_is_broadcasting], - x['times'][b,:t][times_where_this_tracked_is_broadcasting]]) + def __init__( + self, + d_drone_state=3 + 4, + d_radio_feature=258, + d_detector_observation_embedding=128, + d_trajectory_embedding=256, + trajectory_prediction_n_layers=8, + d_trajectory_prediction_output=(2 + 2 + 1) + (2 + 2 + 1), + d_model=512, + n_heads=8, + d_hid=256, + n_layers=4, + n_outputs=8, + ssn_d_hid=64, + ssn_n_layers=8, + # ssn_d_output=5, + ): + super().__init__() + self.d_detector_observation_embedding = d_detector_observation_embedding + self.d_trajectory_embedding = d_trajectory_embedding + self.d_trajectory_prediction_output = d_trajectory_prediction_output + + self.snap_shot_net = SnapShotEmbeddingNet( + d_in=d_radio_feature + d_drone_state, # time is inside drone state + d_hid=ssn_d_hid, + # d_out=5,# 2 mean, 2 sigma, 1 angle + d_embed=d_detector_observation_embedding, + n_layers=ssn_n_layers, ) - #breakpoint() - assert(nested_tensors[-1].shape[0]!=0) - max_snapshots=max(max_snapshots,nested_tensors[-1].shape[0]) # efficiency, what is the biggest we need to compute - idxs.append((b.item(),tracked)) - #breakpoint() - #(Pdb) x['times'].shape - #torch.Size([32, 256, 1]) - - #move all the drone+observation sequences into a common tensor with padding - tformer_input=torch.zeros((len(idxs),max_snapshots,self.d_detector_observation_embedding+1),device=device) - src_key_padding_mask=torch.zeros((len(idxs),max_snapshots),dtype=bool) # TODO ones and mask is faster? - for idx,emebeddings_per_batch_and_tracked in enumerate(nested_tensors): - tracked_time_steps,_=emebeddings_per_batch_and_tracked.shape - #breakpoint() - tformer_input[idx,:tracked_time_steps]=emebeddings_per_batch_and_tracked - src_key_padding_mask[idx,tracked_time_steps:]=True - src_key_padding_mask=src_key_padding_mask.to(device) #TODO initialize on device? - #run the self attention layer # this takes most of the TIME! - self_attention_output=self.tformer(tformer_input,src_key_padding_mask=src_key_padding_mask) - - #select or aggregate the emebedding to use from the historical options - trajectory_embeddings=torch.zeros((batch_size,n_emitters,self.d_trajectory_embedding),device=device) - for idx,(b,emitter_idx) in enumerate(idxs): - #trajectory_embeddings[b,emitter_idx]=self_attention_output[idx,~src_key_padding_mask[idx]].mean(axis=0) - trajectory_embeddings[b,emitter_idx]=self_attention_output[idx,~src_key_padding_mask[idx]][-1] - - ####### - # TRAJECTORY PREDICTIONS - ####### - time_fractions=torch.linspace(0,1,time_steps,device=device)[...,None] # TODO add an assert to make sure this does not go out of sync with other format - trajectory_predictions=torch.zeros((batch_size,n_emitters,time_steps,self.d_trajectory_prediction_output),device=device) - for b in torch.arange(batch_size): - for emitter_idx in torch.arange(n_emitters): - trajectory_input=torch.cat([ - trajectory_embeddings[b,emitter_idx][None].expand((time_steps,self.d_trajectory_embedding)), - time_fractions, - #TODO ADD DRONE EMBEDDING! - ],axis=1) - trajectory_predictions[b,emitter_idx]=self.trajectory_prediction_net(trajectory_input)['output'] - - return { - 'single_snapshot_predictions':single_snapshot_predictions, - #'trajectory_embeddings':trajectory_embeddings, - 'trajectory_predictions':trajectory_predictions.transpose(2,1), - } + self.trajectory_prediction_net = EmbeddingNet( + d_in=d_trajectory_embedding + 1, + d_hid=d_trajectory_embedding, + d_out=d_trajectory_prediction_output, + d_embed=d_trajectory_embedding, + n_layers=trajectory_prediction_n_layers, + ) + + self.tformer = TransformerEncOnlyModel( + d_in=d_detector_observation_embedding + 1, + d_model=d_model, + n_heads=n_heads, + d_hid=d_hid, + n_layers=n_layers, + dropout=0.1, + n_outputs=d_trajectory_embedding, + ) + + def l2(self): + return { + "snap_shot_net": sum( + [(p**2).sum() for p in self.snap_shot_net.parameters()] + ) + } + + def forward(self, x): + # Lets get the detector observation embeddings + batch_size, time_steps, n_emitters, _ = x["emitters_n_broadcasts"].shape + device = x["emitters_n_broadcasts"].device + + ############### + # SINGLE SNAPSHOT PREDICTIONS + EMBEDDINGS + ############### + d = self.snap_shot_net( + torch.cat( + [ # drone state + radio feature as input to SSN + x["drone_state"], + x["radio_feature"], + ], + dim=2, + ) + ) + drone_state_and_observation_embeddings = d["embedding"].detach() + single_snapshot_predictions = d["output"] + + ################ + # TRAJECTORY EMBEDDINGS + # pick random time for each batch item, crop session, and compute tracjectory emebedding for each tracked object + ################ + # generate random times to grab + if self.training: + rt = torch.randint( + low=max(2, batch_size // 10), high=time_steps - 1, size=(batch_size,) + ) # keep on CPU?, low=2 here gaurantees at least one thing was being tracked for each example + rt[: batch_size // 4] = ( + time_steps - 1 + ) // 2 + 1 # this might make it a bit smoother? + rt[ + batch_size // 2 : ((3 * batch_size) // 4) + ] = time_steps # this might make it a bit smoother? + else: + print("EVAL MODE") + rt = torch.ones(batch_size, dtype=int) * (time_steps - 1) + # now lets grab the (nsources,1) vector for each batch example that tells us how many times the object has transmitted previously + tracking = x["emitters_n_broadcasts"][ + torch.arange(batch_size), rt - 1 + ].cpu() # positive values for things already being tracked + + # find out how many objects are being tracked in each batch example + # so that we can pull out their embeddings to calculate trajectory embeddings + nested_tensors = [] + idxs = [] # List of batch idx and idx of object being tracked + max_snapshots = 0 + for b in torch.arange(batch_size): + t = rt[b] + for tracked in torch.where(tracking[b] > 0)[0]: + # get the mask where this emitter is broadcasting in the first t steps + times_where_this_tracked_is_broadcasting = x["emitters_broadcasting"][ + b, :t, tracked, 0 + ].to(bool) + # pull the drone state and observations for these time steps + # nested_tensors.append( + # drone_state_and_observation_embeddings[b,:t][times_where_this_tracked_is_broadcasting] + # ) + nested_tensors.append( + torch.hstack( + [ + drone_state_and_observation_embeddings[b, :t][ + times_where_this_tracked_is_broadcasting + ], + x["times"][b, :t][times_where_this_tracked_is_broadcasting], + ] + ) + ) + # breakpoint() + assert nested_tensors[-1].shape[0] != 0 + max_snapshots = max( + max_snapshots, nested_tensors[-1].shape[0] + ) # efficiency, what is the biggest we need to compute + idxs.append((b.item(), tracked)) + # breakpoint() + # (Pdb) x['times'].shape + # torch.Size([32, 256, 1]) + + # move all the drone+observation sequences into a common tensor with padding + tformer_input = torch.zeros( + (len(idxs), max_snapshots, self.d_detector_observation_embedding + 1), + device=device, + ) + src_key_padding_mask = torch.zeros( + (len(idxs), max_snapshots), dtype=bool + ) # TODO ones and mask is faster? + for idx, emebeddings_per_batch_and_tracked in enumerate(nested_tensors): + tracked_time_steps, _ = emebeddings_per_batch_and_tracked.shape + # breakpoint() + tformer_input[idx, :tracked_time_steps] = emebeddings_per_batch_and_tracked + src_key_padding_mask[idx, tracked_time_steps:] = True + src_key_padding_mask = src_key_padding_mask.to( + device + ) # TODO initialize on device? + # run the self attention layer # this takes most of the TIME! + self_attention_output = self.tformer( + tformer_input, src_key_padding_mask=src_key_padding_mask + ) + + # select or aggregate the emebedding to use from the historical options + trajectory_embeddings = torch.zeros( + (batch_size, n_emitters, self.d_trajectory_embedding), device=device + ) + for idx, (b, emitter_idx) in enumerate(idxs): + # trajectory_embeddings[b,emitter_idx]=self_attention_output[idx,~src_key_padding_mask[idx]].mean(axis=0) + trajectory_embeddings[b, emitter_idx] = self_attention_output[ + idx, ~src_key_padding_mask[idx] + ][-1] + + ####### + # TRAJECTORY PREDICTIONS + ####### + time_fractions = torch.linspace(0, 1, time_steps, device=device)[ + ..., None + ] # TODO add an assert to make sure this does not go out of sync with other format + trajectory_predictions = torch.zeros( + (batch_size, n_emitters, time_steps, self.d_trajectory_prediction_output), + device=device, + ) + for b in torch.arange(batch_size): + for emitter_idx in torch.arange(n_emitters): + trajectory_input = torch.cat( + [ + trajectory_embeddings[b, emitter_idx][None].expand( + (time_steps, self.d_trajectory_embedding) + ), + time_fractions, + # TODO ADD DRONE EMBEDDING! + ], + axis=1, + ) + trajectory_predictions[b, emitter_idx] = self.trajectory_prediction_net( + trajectory_input + )["output"] + + return { + "single_snapshot_predictions": single_snapshot_predictions, + #'trajectory_embeddings':trajectory_embeddings, + "trajectory_predictions": trajectory_predictions.transpose(2, 1), + } class EmbeddingNet(nn.Module): - def __init__(self, - d_in, - d_out, - d_hid, - d_embed, - n_layers): - super(EmbeddingNet,self).__init__() - - self.embed_net=nn.Sequential( - nn.Linear(d_in,d_hid), - nn.SELU(), - *[SkipConnection( - nn.Sequential( - nn.Linear(d_hid,d_hid), - nn.LayerNorm(d_hid), - #nn.ReLU() - nn.SELU() - )) - for _ in range(n_layers) ], - #nn.LayerNorm(d_hid), - nn.Linear(d_hid,d_embed), - nn.LayerNorm(d_embed), - ) - self.lin_output=nn.Linear(d_embed,d_out) - - def forward(self, x): - embed=self.embed_net(x) - output=self.lin_output(embed) - - if output.isnan().any(): - breakpoint() - return {'embedding':embed, - 'output':output} + def __init__(self, d_in, d_out, d_hid, d_embed, n_layers): + super(EmbeddingNet, self).__init__() + + self.embed_net = nn.Sequential( + nn.Linear(d_in, d_hid), + nn.SELU(), + *[ + SkipConnection( + nn.Sequential( + nn.Linear(d_hid, d_hid), + nn.LayerNorm(d_hid), + # nn.ReLU() + nn.SELU(), + ) + ) + for _ in range(n_layers) + ], + # nn.LayerNorm(d_hid), + nn.Linear(d_hid, d_embed), + nn.LayerNorm(d_embed), + ) + self.lin_output = nn.Linear(d_embed, d_out) + + def forward(self, x): + embed = self.embed_net(x) + output = self.lin_output(embed) + + if output.isnan().any(): + breakpoint() + return {"embedding": embed, "output": output} + class SnapShotEmbeddingNet(nn.Module): - def __init__(self, - d_in, - d_hid, - d_embed, - n_layers, - directional=True): - super(SnapShotEmbeddingNet,self).__init__() - self.embedding_net=EmbeddingNet( - d_in=d_in, - d_out=5, # 2 mean , 2 sigma, 1 angle - d_hid=d_hid, - d_embed=d_embed, - n_layers=n_layers) - self.directional=True - - - def forward(self, x): - d=self.embedding_net(x) - if not self.directional: - return d - #else lets make sure the angle is the angle between mean and drone - #so that we get gaussians along the angle - - diff=d['output'][:,:,:2]-x[:,:,:2] # vector from detector to predicted mean - d['output'][:,:,4]=torch.arctan2(diff[...,1],diff[...,0]) - #print("DIFF",diff[0,0,:]) - #print("THETA",d['output'][0,0,4]) - #d['output'][:,:,3]=0.5 - #d['output'][:,:,2]=5 - return d + def __init__(self, d_in, d_hid, d_embed, n_layers, directional=True): + super(SnapShotEmbeddingNet, self).__init__() + self.embedding_net = EmbeddingNet( + d_in=d_in, + d_out=5, # 2 mean , 2 sigma, 1 angle + d_hid=d_hid, + d_embed=d_embed, + n_layers=n_layers, + ) + self.directional = True + + def forward(self, x): + d = self.embedding_net(x) + if not self.directional: + return d + # else lets make sure the angle is the angle between mean and drone + # so that we get gaussians along the angle + + diff = ( + d["output"][:, :, :2] - x[:, :, :2] + ) # vector from detector to predicted mean + d["output"][:, :, 4] = torch.arctan2(diff[..., 1], diff[..., 0]) + # print("DIFF",diff[0,0,:]) + # print("THETA",d['output'][0,0,4]) + # d['output'][:,:,3]=0.5 + # d['output'][:,:,2]=5 + return d class SnapshotNet(nn.Module): - def __init__(self, - snapshots_per_sample=1, - d_drone_state=4+4, - d_radio_feature=258, - d_model=512, - n_heads=8, - d_hid=256, - n_layers=1, - n_outputs=8, - dropout=0.0, - ssn_d_hid=64, - ssn_n_layers=8, - ssn_n_outputs=8, - ssn_d_embed=64, - ssn_dropout=0.0, - tformer_input=['drone_state','radio_feature','embedding','single_snapshot_pred'], - positional_encoding_len=0): - super().__init__() - self.d_radio_feature=d_radio_feature - self.d_drone_state=d_drone_state - self.d_model=d_model - self.n_heads=n_heads - self.d_hid=d_hid - self.n_outputs=n_outputs - self.dropout=dropout - self.tformer_input=tformer_input - self.tformer_input_dim=0 - for k,v in [ - ('drone_state',d_drone_state), - ('radio_feature',d_radio_feature), - ('embedding',ssn_d_embed), - ('single_snapshot_pred',n_outputs)]: - if k in self.tformer_input: - self.tformer_input_dim+=v - self.tformer_input_dim+=positional_encoding_len - - - self.snap_shot_net=SingleSnapshotNet( - d_input_feature=d_radio_feature+d_drone_state, - d_hid=ssn_d_hid, - d_embed=ssn_d_embed, - n_layers=ssn_n_layers, - n_outputs=ssn_n_outputs, - dropout=ssn_dropout) - #self.snap_shot_net=Task1Net(d_radio_feature*snapshots_per_sample) - - self.tformer=TransformerEncOnlyModel( - d_in=self.tformer_input_dim, - #d_radio_feature=ssn_d_embed, #+n_outputs, - d_model=d_model, - n_heads=n_heads, - d_hid=d_hid, - n_layers=n_layers, - dropout=dropout, - n_outputs=n_outputs) - - self.positional_encoding=nn.Identity() - if positional_encoding_len>0: - self.positional_encoding=PositionalEncoding(positional_encoding_len,snapshots_per_sample) - - def forward(self,x): - d=self.snap_shot_net(x) - #return single_snapshot_output,single_snapshot_output - tformer_input=self.positional_encoding(torch.cat([d[t] for t in self.tformer_input ],axis=2)) - tformer_output=self.tformer(tformer_input) - return {'transformer_pred':tformer_output, - 'single_snapshot_pred':d['single_snapshot_pred']} + def __init__( + self, + snapshots_per_sample=1, + d_drone_state=4 + 4, + d_radio_feature=258, + d_model=512, + n_heads=8, + d_hid=256, + n_layers=1, + n_outputs=8, + dropout=0.0, + ssn_d_hid=64, + ssn_n_layers=8, + ssn_n_outputs=8, + ssn_d_embed=64, + ssn_dropout=0.0, + tformer_input=[ + "drone_state", + "radio_feature", + "embedding", + "single_snapshot_pred", + ], + positional_encoding_len=0, + ): + super().__init__() + self.d_radio_feature = d_radio_feature + self.d_drone_state = d_drone_state + self.d_model = d_model + self.n_heads = n_heads + self.d_hid = d_hid + self.n_outputs = n_outputs + self.dropout = dropout + self.tformer_input = tformer_input + self.tformer_input_dim = 0 + for k, v in [ + ("drone_state", d_drone_state), + ("radio_feature", d_radio_feature), + ("embedding", ssn_d_embed), + ("single_snapshot_pred", n_outputs), + ]: + if k in self.tformer_input: + self.tformer_input_dim += v + self.tformer_input_dim += positional_encoding_len + + self.snap_shot_net = SingleSnapshotNet( + d_input_feature=d_radio_feature + d_drone_state, + d_hid=ssn_d_hid, + d_embed=ssn_d_embed, + n_layers=ssn_n_layers, + n_outputs=ssn_n_outputs, + dropout=ssn_dropout, + ) + # self.snap_shot_net=Task1Net(d_radio_feature*snapshots_per_sample) + + self.tformer = TransformerEncOnlyModel( + d_in=self.tformer_input_dim, + # d_radio_feature=ssn_d_embed, #+n_outputs, + d_model=d_model, + n_heads=n_heads, + d_hid=d_hid, + n_layers=n_layers, + dropout=dropout, + n_outputs=n_outputs, + ) + + self.positional_encoding = nn.Identity() + if positional_encoding_len > 0: + self.positional_encoding = PositionalEncoding( + positional_encoding_len, snapshots_per_sample + ) + + def forward(self, x): + d = self.snap_shot_net(x) + # return single_snapshot_output,single_snapshot_output + tformer_input = self.positional_encoding( + torch.cat([d[t] for t in self.tformer_input], axis=2) + ) + tformer_output = self.tformer(tformer_input) + return { + "transformer_pred": tformer_output, + "single_snapshot_pred": d["single_snapshot_pred"], + } + class SingleSnapshotNet(nn.Module): - def __init__(self, - d_input_feature, - d_hid, - d_embed, - n_layers, - n_outputs, - dropout, - snapshots_per_sample=0): - super(SingleSnapshotNet,self).__init__() - self.snapshots_per_sample=snapshots_per_sample - self.d_input_feature=d_input_feature - if self.snapshots_per_sample>0: - self.d_input_feature*=snapshots_per_sample - self.d_hid=d_hid - self.d_embed=d_embed - self.n_layers=n_layers - self.n_outputs=n_outputs - self.dropout=dropout - - self.embed_net=nn.Sequential( - #nn.LayerNorm(self.d_radio_feature), - nn.Linear(self.d_input_feature,d_hid), - nn.SELU(), - *[SkipConnection( - nn.Sequential( - nn.LayerNorm(d_hid), - nn.Linear(d_hid,d_hid), - #nn.ReLU() - nn.SELU() - )) - for _ in range(n_layers) ], - #nn.LayerNorm(d_hid), - nn.Linear(d_hid,d_embed), - nn.LayerNorm(d_embed) - ) - self.lin_output=nn.Linear(d_embed,self.n_outputs) - - def forward(self, x_dict): - x=torch.cat([x_dict['drone_state'],x_dict['radio_feature']],axis=2) - if self.snapshots_per_sample>0: - x=x.reshape(x.shape[0],-1) - embed=self.embed_net(x) - output=self.lin_output(embed) - if output.isnan().any(): - breakpoint() - if self.snapshots_per_sample>0: - output=output.reshape(-1,1,self.n_outputs) - return {'fc_pred':output} - return {'drone_state':x_dict['drone_state'], - 'radio_feature':x_dict['radio_feature'], - 'single_snapshot_pred':output,'embedding':embed} + def __init__( + self, + d_input_feature, + d_hid, + d_embed, + n_layers, + n_outputs, + dropout, + snapshots_per_sample=0, + ): + super(SingleSnapshotNet, self).__init__() + self.snapshots_per_sample = snapshots_per_sample + self.d_input_feature = d_input_feature + if self.snapshots_per_sample > 0: + self.d_input_feature *= snapshots_per_sample + self.d_hid = d_hid + self.d_embed = d_embed + self.n_layers = n_layers + self.n_outputs = n_outputs + self.dropout = dropout + + self.embed_net = nn.Sequential( + # nn.LayerNorm(self.d_radio_feature), + nn.Linear(self.d_input_feature, d_hid), + nn.SELU(), + *[ + SkipConnection( + nn.Sequential( + nn.LayerNorm(d_hid), + nn.Linear(d_hid, d_hid), + # nn.ReLU() + nn.SELU(), + ) + ) + for _ in range(n_layers) + ], + # nn.LayerNorm(d_hid), + nn.Linear(d_hid, d_embed), + nn.LayerNorm(d_embed), + ) + self.lin_output = nn.Linear(d_embed, self.n_outputs) + + def forward(self, x_dict): + x = torch.cat([x_dict["drone_state"], x_dict["radio_feature"]], axis=2) + if self.snapshots_per_sample > 0: + x = x.reshape(x.shape[0], -1) + embed = self.embed_net(x) + output = self.lin_output(embed) + if output.isnan().any(): + breakpoint() + if self.snapshots_per_sample > 0: + output = output.reshape(-1, 1, self.n_outputs) + return {"fc_pred": output} + return { + "drone_state": x_dict["drone_state"], + "radio_feature": x_dict["radio_feature"], + "single_snapshot_pred": output, + "embedding": embed, + } + class Task1Net(nn.Module): - def __init__(self,ndim,n_outputs=8): - super().__init__() - self.bn1 = nn.BatchNorm1d(120) - self.bn2 = nn.BatchNorm1d(84) - self.bn3 = nn.BatchNorm1d(n_outputs) - self.fc1 = nn.Linear(ndim, 120) - self.fc2 = nn.Linear(120, 84) - self.fc3 = nn.Linear(84, n_outputs) - self.n_outputs=n_outputs - self.ndim=ndim - - def forward(self, x): - x = x.reshape(x.shape[0],-1) - x = F.relu(self.bn1(self.fc1(x))) - x = F.relu(self.bn2(self.fc2(x))) - x = F.relu(self.bn3(self.fc3(x))) - x = x.reshape(-1,1,self.n_outputs) - return {'fc_pred':x} #.reshape(x.shape[0],1,2) + def __init__(self, ndim, n_outputs=8): + super().__init__() + self.bn1 = nn.BatchNorm1d(120) + self.bn2 = nn.BatchNorm1d(84) + self.bn3 = nn.BatchNorm1d(n_outputs) + self.fc1 = nn.Linear(ndim, 120) + self.fc2 = nn.Linear(120, 84) + self.fc3 = nn.Linear(84, n_outputs) + self.n_outputs = n_outputs + self.ndim = ndim + + def forward(self, x): + x = x.reshape(x.shape[0], -1) + x = F.relu(self.bn1(self.fc1(x))) + x = F.relu(self.bn2(self.fc2(x))) + x = F.relu(self.bn3(self.fc3(x))) + x = x.reshape(-1, 1, self.n_outputs) + return {"fc_pred": x} # .reshape(x.shape[0],1,2) class UNet(nn.Module): + def __init__(self, in_channels=3, out_channels=1, width=128, init_features=32): + super(UNet, self).__init__() + + features = init_features + self.encoder1 = UNet._block(in_channels, features, name="enc1") + self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2) + self.encoder2 = UNet._block(features, features * 2, name="enc2") + self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2) + self.encoder3 = UNet._block(features * 2, features * 4, name="enc3") + self.pool3 = nn.MaxPool2d(kernel_size=2, stride=2) + self.encoder4 = UNet._block(features * 4, features * 8, name="enc4") + self.pool4 = nn.MaxPool2d(kernel_size=2, stride=2) + self.encoder5 = UNet._block(features * 8, features * 16, name="enc4") + self.pool5 = nn.MaxPool2d(kernel_size=2, stride=2) + + self.bottleneck = UNet._block(features * 16, features * 32, name="bottleneck") + # self.bottleneck = UNet._block(features * 8, features * 16, name="bottleneck") + + self.upconv5 = nn.ConvTranspose2d( + features * 32, features * 16, kernel_size=2, stride=2 + ) + self.decoder5 = UNet._block((features * 16) * 2, features * 16, name="dec5") + self.upconv4 = nn.ConvTranspose2d( + features * 16, features * 8, kernel_size=2, stride=2 + ) + self.decoder4 = UNet._block((features * 8) * 2, features * 8, name="dec4") + self.upconv3 = nn.ConvTranspose2d( + features * 8, features * 4, kernel_size=2, stride=2 + ) + self.decoder3 = UNet._block((features * 4) * 2, features * 4, name="dec3") + self.upconv2 = nn.ConvTranspose2d( + features * 4, features * 2, kernel_size=2, stride=2 + ) + self.decoder2 = UNet._block((features * 2) * 2, features * 2, name="dec2") + self.upconv1 = nn.ConvTranspose2d( + features * 2, features, kernel_size=2, stride=2 + ) + self.decoder1 = UNet._block(features * 2, features, name="dec1") - def __init__(self, in_channels=3, out_channels=1, width=128, init_features=32): - super(UNet, self).__init__() - - features = init_features - self.encoder1 = UNet._block(in_channels, features, name="enc1") - self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2) - self.encoder2 = UNet._block(features, features * 2, name="enc2") - self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2) - self.encoder3 = UNet._block(features * 2, features * 4, name="enc3") - self.pool3 = nn.MaxPool2d(kernel_size=2, stride=2) - self.encoder4 = UNet._block(features * 4, features * 8, name="enc4") - self.pool4 = nn.MaxPool2d(kernel_size=2, stride=2) - self.encoder5 = UNet._block(features * 8, features * 16, name="enc4") - self.pool5 = nn.MaxPool2d(kernel_size=2, stride=2) - - self.bottleneck = UNet._block(features * 16, features * 32, name="bottleneck") - #self.bottleneck = UNet._block(features * 8, features * 16, name="bottleneck") - - self.upconv5 = nn.ConvTranspose2d( - features * 32, features * 16, kernel_size=2, stride=2 - ) - self.decoder5 = UNet._block((features * 16) * 2, features * 16, name="dec5") - self.upconv4 = nn.ConvTranspose2d( - features * 16, features * 8, kernel_size=2, stride=2 - ) - self.decoder4 = UNet._block((features * 8) * 2, features * 8, name="dec4") - self.upconv3 = nn.ConvTranspose2d( - features * 8, features * 4, kernel_size=2, stride=2 - ) - self.decoder3 = UNet._block((features * 4) * 2, features * 4, name="dec3") - self.upconv2 = nn.ConvTranspose2d( - features * 4, features * 2, kernel_size=2, stride=2 - ) - self.decoder2 = UNet._block((features * 2) * 2, features * 2, name="dec2") - self.upconv1 = nn.ConvTranspose2d( - features * 2, features, kernel_size=2, stride=2 - ) - self.decoder1 = UNet._block(features * 2, features, name="dec1") - - self.conv = nn.Conv2d( - in_channels=features, out_channels=out_channels, kernel_size=1 - ) - - def forward(self, x): - enc1 = self.encoder1(x) - enc2 = self.encoder2(self.pool1(enc1)) - enc3 = self.encoder3(self.pool2(enc2)) - enc4 = self.encoder4(self.pool3(enc3)) - enc5 = self.encoder5(self.pool4(enc4)) - - bottleneck = self.bottleneck(self.pool5(enc5)) - #bottleneck = self.bottleneck(self.pool4(enc4)) - - dec5 = self.upconv5(bottleneck) - dec5 = torch.cat((dec5, enc5), dim=1) - dec5 = self.decoder5(dec5) - dec4 = self.upconv4(dec5) - dec4 = torch.cat((dec4, enc4), dim=1) - dec4 = self.decoder4(dec4) - dec3 = self.upconv3(dec4) - dec3 = torch.cat((dec3, enc3), dim=1) - dec3 = self.decoder3(dec3) - dec2 = self.upconv2(dec3) - dec2 = torch.cat((dec2, enc2), dim=1) - dec2 = self.decoder2(dec2) - dec1 = self.upconv1(dec2) - dec1 = torch.cat((dec1, enc1), dim=1) - dec1 = self.decoder1(dec1) - out=torch.sigmoid(self.conv(dec1)) - return {'image_preds':out/out.sum(axis=[2,3],keepdims=True)} - - @staticmethod - def _block(in_channels, features, name): - return nn.Sequential( - OrderedDict( - [ - ( - name + "conv1", - nn.Conv2d( - in_channels=in_channels, - out_channels=features, - kernel_size=3, - padding=1, - bias=False, - ), - ), - (name + "norm1", nn.BatchNorm2d(num_features=features)), - (name + "relu1", nn.ReLU(inplace=True)), - ( - name + "conv2", - nn.Conv2d( - in_channels=features, - out_channels=features, - kernel_size=3, - padding=1, - bias=False, - ), - ), - (name + "norm2", nn.BatchNorm2d(num_features=features)), - (name + "relu2", nn.ReLU(inplace=True)), - ] - ) - ) + self.conv = nn.Conv2d( + in_channels=features, out_channels=out_channels, kernel_size=1 + ) + + def forward(self, x): + enc1 = self.encoder1(x) + enc2 = self.encoder2(self.pool1(enc1)) + enc3 = self.encoder3(self.pool2(enc2)) + enc4 = self.encoder4(self.pool3(enc3)) + enc5 = self.encoder5(self.pool4(enc4)) + + bottleneck = self.bottleneck(self.pool5(enc5)) + # bottleneck = self.bottleneck(self.pool4(enc4)) + + dec5 = self.upconv5(bottleneck) + dec5 = torch.cat((dec5, enc5), dim=1) + dec5 = self.decoder5(dec5) + dec4 = self.upconv4(dec5) + dec4 = torch.cat((dec4, enc4), dim=1) + dec4 = self.decoder4(dec4) + dec3 = self.upconv3(dec4) + dec3 = torch.cat((dec3, enc3), dim=1) + dec3 = self.decoder3(dec3) + dec2 = self.upconv2(dec3) + dec2 = torch.cat((dec2, enc2), dim=1) + dec2 = self.decoder2(dec2) + dec1 = self.upconv1(dec2) + dec1 = torch.cat((dec1, enc1), dim=1) + dec1 = self.decoder1(dec1) + out = torch.sigmoid(self.conv(dec1)) + return {"image_preds": out / out.sum(axis=[2, 3], keepdims=True)} + + @staticmethod + def _block(in_channels, features, name): + return nn.Sequential( + OrderedDict( + [ + ( + name + "conv1", + nn.Conv2d( + in_channels=in_channels, + out_channels=features, + kernel_size=3, + padding=1, + bias=False, + ), + ), + (name + "norm1", nn.BatchNorm2d(num_features=features)), + (name + "relu1", nn.ReLU(inplace=True)), + ( + name + "conv2", + nn.Conv2d( + in_channels=features, + out_channels=features, + kernel_size=3, + padding=1, + bias=False, + ), + ), + (name + "norm2", nn.BatchNorm2d(num_features=features)), + (name + "relu2", nn.ReLU(inplace=True)), + ] + ) + ) diff --git a/software/model_training_and_inference/test.py b/software/model_training_and_inference/test.py index b3b903ab..6f97b90f 100644 --- a/software/model_training_and_inference/test.py +++ b/software/model_training_and_inference/test.py @@ -59,12 +59,19 @@ from torch.nn import TransformerEncoder, TransformerEncoderLayer from torch.utils.data import dataset -class TransformerModel(nn.Module): - def __init__(self, ntoken: int, d_model: int, nhead: int, d_hid: int, - nlayers: int, dropout: float = 0.5): +class TransformerModel(nn.Module): + def __init__( + self, + ntoken: int, + d_model: int, + nhead: int, + d_hid: int, + nlayers: int, + dropout: float = 0.5, + ): super().__init__() - self.model_type = 'Transformer' + self.model_type = "Transformer" self.pos_encoder = PositionalEncoding(d_model, dropout) encoder_layers = TransformerEncoderLayer(d_model, nhead, d_hid, dropout) self.transformer_encoder = TransformerEncoder(encoder_layers, nlayers) @@ -104,25 +111,27 @@ def forward(self, src: Tensor, src_mask: Tensor = None) -> Tensor: # different frequencies. # -class PositionalEncoding(nn.Module): +class PositionalEncoding(nn.Module): def __init__(self, d_model: int, dropout: float = 0.1, max_len: int = 5000): super().__init__() self.dropout = nn.Dropout(p=dropout) position = torch.arange(max_len).unsqueeze(1) - div_term = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model)) + div_term = torch.exp( + torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model) + ) pe = torch.zeros(max_len, 1, d_model) pe[:, 0, 0::2] = torch.sin(position * div_term) pe[:, 0, 1::2] = torch.cos(position * div_term) - self.register_buffer('pe', pe) + self.register_buffer("pe", pe) def forward(self, x: Tensor) -> Tensor: """ Arguments: x: Tensor, shape ``[seq_len, batch_size, embedding_dim]`` """ - x = x + self.pe[:x.size(0)] + x = x + self.pe[: x.size(0)] return self.dropout(x) @@ -172,16 +181,20 @@ def forward(self, x: Tensor) -> Tensor: from torchtext.data.utils import get_tokenizer from torchtext.vocab import build_vocab_from_iterator -train_iter = WikiText2(split='train') -tokenizer = get_tokenizer('basic_english') -vocab = build_vocab_from_iterator(map(tokenizer, train_iter), specials=['']) -vocab.set_default_index(vocab['']) +train_iter = WikiText2(split="train") +tokenizer = get_tokenizer("basic_english") +vocab = build_vocab_from_iterator(map(tokenizer, train_iter), specials=[""]) +vocab.set_default_index(vocab[""]) + def data_process(raw_text_iter: dataset.IterableDataset) -> Tensor: """Converts raw text into a flat Tensor.""" - data = [torch.tensor(vocab(tokenizer(item)), dtype=torch.long) for item in raw_text_iter] + data = [ + torch.tensor(vocab(tokenizer(item)), dtype=torch.long) for item in raw_text_iter + ] return torch.cat(tuple(filter(lambda t: t.numel() > 0, data))) + # ``train_iter`` was "consumed" by the process of building the vocab, # so we have to create it again train_iter, val_iter, test_iter = WikiText2() @@ -189,7 +202,8 @@ def data_process(raw_text_iter: dataset.IterableDataset) -> Tensor: val_data = data_process(val_iter) test_data = data_process(test_iter) -device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') +device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + def batchify(data: Tensor, bsz: int) -> Tensor: """Divides the data into ``bsz`` separate sequences, removing extra elements @@ -203,10 +217,11 @@ def batchify(data: Tensor, bsz: int) -> Tensor: Tensor of shape ``[N // bsz, bsz]`` """ seq_len = data.size(0) // bsz - data = data[:seq_len * bsz] + data = data[: seq_len * bsz] data = data.view(bsz, seq_len).t().contiguous() return data.to(device) + batch_size = 20 eval_batch_size = 10 train_data = batchify(train_data, batch_size) # shape ``[seq_len, batch_size]`` @@ -235,6 +250,8 @@ def batchify(data: Tensor, bsz: int) -> Tensor: # bptt = 35 + + def get_batch(source: Tensor, i: int) -> Tuple[Tensor, Tensor]: """ Args: @@ -246,8 +263,8 @@ def get_batch(source: Tensor, i: int) -> Tuple[Tensor, Tensor]: target has shape ``[seq_len * batch_size]`` """ seq_len = min(bptt, len(source) - 1 - i) - data = source[i:i+seq_len] - target = source[i+1:i+1+seq_len].reshape(-1) + data = source[i : i + seq_len] + target = source[i + 1 : i + 1 + seq_len].reshape(-1) return data, target @@ -293,9 +310,10 @@ def get_batch(source: Tensor, i: int) -> Tuple[Tensor, Tensor]: optimizer = torch.optim.SGD(model.parameters(), lr=lr) scheduler = torch.optim.lr_scheduler.StepLR(optimizer, 1.0, gamma=0.95) + def train(model: nn.Module) -> None: model.train() # turn on train mode - total_loss = 0. + total_loss = 0.0 log_interval = 200 start_time = time.time() @@ -318,15 +336,18 @@ def train(model: nn.Module) -> None: ms_per_batch = (time.time() - start_time) * 1000 / log_interval cur_loss = total_loss / log_interval ppl = math.exp(cur_loss) - print(f'| epoch {epoch:3d} | {batch:5d}/{num_batches:5d} batches | ' - f'lr {lr:02.2f} | ms/batch {ms_per_batch:5.2f} | ' - f'loss {cur_loss:5.2f} | ppl {ppl:8.2f}') + print( + f"| epoch {epoch:3d} | {batch:5d}/{num_batches:5d} batches | " + f"lr {lr:02.2f} | ms/batch {ms_per_batch:5.2f} | " + f"loss {cur_loss:5.2f} | ppl {ppl:8.2f}" + ) total_loss = 0 start_time = time.time() + def evaluate(model: nn.Module, eval_data: Tensor) -> float: model.eval() # turn on evaluation mode - total_loss = 0. + total_loss = 0.0 with torch.no_grad(): for i in range(0, eval_data.size(0) - 1, bptt): data, targets = get_batch(eval_data, i) @@ -336,11 +357,12 @@ def evaluate(model: nn.Module, eval_data: Tensor) -> float: total_loss += seq_len * criterion(output_flat, targets).item() return total_loss / (len(eval_data) - 1) + ###################################################################### # Loop over epochs. Save the model if the validation loss is the best # we've seen so far. Adjust the learning rate after each epoch. -best_val_loss = float('inf') +best_val_loss = float("inf") epochs = 3 with TemporaryDirectory() as tempdir: @@ -352,17 +374,19 @@ def evaluate(model: nn.Module, eval_data: Tensor) -> float: val_loss = evaluate(model, val_data) val_ppl = math.exp(val_loss) elapsed = time.time() - epoch_start_time - print('-' * 89) - print(f'| end of epoch {epoch:3d} | time: {elapsed:5.2f}s | ' - f'valid loss {val_loss:5.2f} | valid ppl {val_ppl:8.2f}') - print('-' * 89) + print("-" * 89) + print( + f"| end of epoch {epoch:3d} | time: {elapsed:5.2f}s | " + f"valid loss {val_loss:5.2f} | valid ppl {val_ppl:8.2f}" + ) + print("-" * 89) if val_loss < best_val_loss: best_val_loss = val_loss torch.save(model.state_dict(), best_model_params_path) scheduler.step() - model.load_state_dict(torch.load(best_model_params_path)) # load best model states + model.load_state_dict(torch.load(best_model_params_path)) # load best model states ###################################################################### @@ -372,7 +396,6 @@ def evaluate(model: nn.Module, eval_data: Tensor) -> float: test_loss = evaluate(model, test_data) test_ppl = math.exp(test_loss) -print('=' * 89) -print(f'| End of training | test loss {test_loss:5.2f} | ' - f'test ppl {test_ppl:8.2f}') -print('=' * 89) +print("=" * 89) +print(f"| End of training | test loss {test_loss:5.2f} | " f"test ppl {test_ppl:8.2f}") +print("=" * 89) diff --git a/software/model_training_and_inference/tests/test_beamformer.py b/software/model_training_and_inference/tests/test_beamformer.py index 03635f74..1cabae81 100644 --- a/software/model_training_and_inference/tests/test_beamformer.py +++ b/software/model_training_and_inference/tests/test_beamformer.py @@ -2,14 +2,14 @@ def test_beamformer(): - #carrier_frequency=2.4e9 - carrier_frequency=100e3 - wavelength=c/carrier_frequency - sampling_frequency=10e6 + # carrier_frequency=2.4e9 + carrier_frequency = 100e3 + wavelength = c / carrier_frequency + sampling_frequency = 10e6 import matplotlib.pyplot as plt - ''' + """ X (source_pos) | | @@ -18,48 +18,58 @@ def test_beamformer(): ---------------- Lets rotate on the source around to the left(counter clockwise) by theta radians - ''' + """ - source_pos=np.array([[0,10000]]) - rotations=[-np.pi/2,-np.pi,-np.pi/4,0,np.pi/2,np.pi/2,np.pi]+list(np.random.uniform(-np.pi,np.pi,10)) - spacings=[wavelength/4] #,wavelength/2,wavelength,wavelength/3]+list(np.random.uniform(0,1,10)) + source_pos = np.array([[0, 10000]]) + rotations = [-np.pi / 2, -np.pi, -np.pi / 4, 0, np.pi / 2, np.pi / 2, np.pi] + list( + np.random.uniform(-np.pi, np.pi, 10) + ) + spacings = [ + wavelength / 4 + ] # ,wavelength/2,wavelength,wavelength/3]+list(np.random.uniform(0,1,10)) - nelements=[2,3,4]#,16] + nelements = [2, 3, 4] # ,16] np.random.seed(1442) - for Detector in [ULADetector,UCADetector]: - for nelement in nelements: - for spacing in spacings: - d=Detector(sampling_frequency,nelement,spacing,sigma=0.0) + for Detector in [ULADetector, UCADetector]: + for nelement in nelements: + for spacing in spacings: + d = Detector(sampling_frequency, nelement, spacing, sigma=0.0) - for rot_theta in rotations: - rot_mat=rotation_matrix(rot_theta) - _source_pos=(rot_mat @ source_pos.T).T # rotate right by rot_theta + for rot_theta in rotations: + rot_mat = rotation_matrix(rot_theta) + _source_pos = ( + rot_mat @ source_pos.T + ).T # rotate right by rot_theta - d.rm_sources() - d.add_source( - IQSource( - _source_pos, # x, y position - carrier_frequency,100e3)) + d.rm_sources() + d.add_source( + IQSource(_source_pos, carrier_frequency, 100e3) # x, y position + ) - signal_matrix=d.get_signal_matrix( - start_time=100, - duration=3/d.sampling_frequency) + signal_matrix = d.get_signal_matrix( + start_time=100, duration=3 / d.sampling_frequency + ) - thetas_at_t,beam_former_outputs_at_t,_=beamformer( - d.all_receiver_pos(), - signal_matrix, - carrier_frequency,spacing=1024+1) + thetas_at_t, beam_former_outputs_at_t, _ = beamformer( + d.all_receiver_pos(), + signal_matrix, + carrier_frequency, + spacing=1024 + 1, + ) - closest_angle_answer=np.argmin(np.abs(thetas_at_t-rot_theta)) - potential_error=np.abs(beam_former_outputs_at_t[:-1]-beam_former_outputs_at_t[1:]).max() - - assert((beam_former_outputs_at_t.max()-potential_error)<=beam_former_outputs_at_t[closest_angle_answer]) + closest_angle_answer = np.argmin(np.abs(thetas_at_t - rot_theta)) + potential_error = np.abs( + beam_former_outputs_at_t[:-1] - beam_former_outputs_at_t[1:] + ).max() - #plt.figure() - #plt.plot(thetas_at_t,beam_former_outputs_at_t) - #plt.axvline(x=-rot_theta) - #plt.axhline(y=beam_former_outputs_at_t.max()-potential_error) - #plt.show() - print("PASS!") + assert ( + beam_former_outputs_at_t.max() - potential_error + ) <= beam_former_outputs_at_t[closest_angle_answer] + # plt.figure() + # plt.plot(thetas_at_t,beam_former_outputs_at_t) + # plt.axvline(x=-rot_theta) + # plt.axhline(y=beam_former_outputs_at_t.max()-potential_error) + # plt.show() + print("PASS!") diff --git a/software/model_training_and_inference/utils/baseline_algorithm.py b/software/model_training_and_inference/utils/baseline_algorithm.py index 7dd61fec..110f7e09 100644 --- a/software/model_training_and_inference/utils/baseline_algorithm.py +++ b/software/model_training_and_inference/utils/baseline_algorithm.py @@ -3,170 +3,209 @@ import torch from PIL import Image import skimage -from utils.image_utils import (detector_positions_to_theta_grid, - labels_to_source_images, radio_to_image) - -def get_top_n_peaks(bf,n=2,peak_size=15): - bf=np.copy(bf) - peaks=np.zeros(n,dtype=int) - for peak_idx in range(n): - #find a peak - peak=bf.argmax() - for idx in np.arange(-peak_size//2,peak_size//2+1): - bf[(peak+idx)%bf.shape[0]]=-np.inf - peaks[peak_idx]=peak - return peaks +from utils.image_utils import ( + detector_positions_to_theta_grid, + labels_to_source_images, + radio_to_image, +) + + +def get_top_n_peaks(bf, n=2, peak_size=15): + bf = np.copy(bf) + peaks = np.zeros(n, dtype=int) + for peak_idx in range(n): + # find a peak + peak = bf.argmax() + for idx in np.arange(-peak_size // 2, peak_size // 2 + 1): + bf[(peak + idx) % bf.shape[0]] = -np.inf + peaks[peak_idx] = peak + return peaks + def frac_to_theta(fracs): - return 2*(fracs-0.5)*np.pi + return 2 * (fracs - 0.5) * np.pi + def line_to_mx(line): - (x,y),((m_x,m_y),)=line - m=m_y/(m_x+1e-9) - return y-m*x,m,(x,y),(m_x,m_y) + (x, y), ((m_x, m_y),) = line + m = m_y / (m_x + 1e-9) + return y - m * x, m, (x, y), (m_x, m_y) + + +def baseline_algorithm(session, width, steps=-1): + line_representations = [] + if steps == -1: + session["beam_former_outputs_at_t"].shape[0] + for idx in np.arange(steps): + _thetas = session["thetas_at_t"][idx][ + get_top_n_peaks(session["beam_former_outputs_at_t"][idx]) + ] + for _theta in _thetas: + direction = np.stack( + [ + np.sin(session["detector_orientation_at_t"][idx] + _theta), + np.cos(session["detector_orientation_at_t"][idx] + _theta), + ], + axis=1, + ) + line_representations.append( + (session["detector_position_at_t"][idx], direction) + ) + return get_top_n_points(line_representations, n=4, width=width, threshold=3) -def baseline_algorithm(session,width,steps=-1): - line_representations=[] - if steps==-1: - session['beam_former_outputs_at_t'].shape[0] - for idx in np.arange(steps): - _thetas=session['thetas_at_t'][idx][get_top_n_peaks(session['beam_former_outputs_at_t'][idx])] - for _theta in _thetas: - direction=np.stack([ - np.sin(session['detector_orientation_at_t'][idx]+_theta), - np.cos(session['detector_orientation_at_t'][idx]+_theta) - ],axis=1) - line_representations.append((session['detector_position_at_t'][idx],direction)) - return get_top_n_points(line_representations,n=4,width=width,threshold=3) +def get_top_n_points( + line_representations, n, width, threshold=3, mn=4, lines_per_step=2 +): + final_points = [] + line_to_point_assignments = np.zeros(len(line_representations), dtype=int) - 1 + line_to_point_distances = np.zeros(len(line_representations), dtype=int) - 1 + for point_idx in range(n): + img = np.zeros((width + 1, width + 1)) + for line_idx in np.arange(len(line_representations)): + if ( + line_to_point_assignments[line_idx] == -1 + or line_to_point_distances[line_idx] >= threshold + ): + line = line_representations[line_idx] + if len(final_points) > 0: # check if its close to the last line + b, m, point, mvec = line_to_mx(line) + fp = final_points[-1] + d = (fp[0] - point[0], fp[1] - point[1]) + if np.sign(mvec[0]) == np.sign(d[0]) and np.sign( + mvec[1] + ) == np.sign( + d[1] + ): # check if its on the right side + mvec_orthogonal = (-mvec[1], mvec[0]) + distance_to_line = abs( + np.dot(d, mvec_orthogonal) / np.linalg.norm(mvec) + ) + if ( + line_to_point_assignments[line_idx] == -1 + or distance_to_line < line_to_point_distances[line_idx] + ): + line_to_point_assignments[line_idx] = len(final_points) - 1 + line_to_point_distances[line_idx] = distance_to_line + if ( + line_to_point_assignments[line_idx] == -1 + or line_to_point_distances[line_idx] >= threshold + ): + b, m, point, mvec = line_to_mx(line) + bp = boundary_point(b, m, point, width, mvec) + rr, cc = skimage.draw.line( + int(point[0]), int(point[1]), int(bp[0]), int(bp[1]) + ) + img[rr, cc] += 1 + if img.max() >= mn: + idx = img.argmax() + p = (idx // (width + 1), idx % (width + 1)) + final_points.append(p) + final_points = np.array(final_points) + points_per_line = ( + np.zeros((len(line_representations) // lines_per_step, 2)) + width / 2 + ) + for output_idx in np.arange(len(line_representations) // lines_per_step): + line_idx = output_idx * lines_per_step + closest_idx = -1 + for _line_idx in range(lines_per_step): + if line_to_point_assignments[line_idx + _line_idx] != -1 and ( + closest_idx == -1 + or line_to_point_distances[line_idx + closest_idx] + > line_to_point_distances[line_idx + _line_idx] + ): + closest_idx = _line_idx + if closest_idx != -1: + points_per_line[output_idx] = final_points[ + line_to_point_assignments[line_idx + closest_idx] + ] + imgs = np.zeros((n, width + 1, width + 1)) + for point_idx in range(n): + for line_idx in np.arange(len(line_representations)): + if line_to_point_assignments[line_idx] >= 0: + line = line_representations[line_idx] + if line_to_point_assignments[line_idx] >= 0: + b, m, point, mvec = line_to_mx(line) + bp = boundary_point(b, m, point, width, mvec) + rr, cc = skimage.draw.line( + int(point[0]), int(point[1]), int(bp[0]), int(bp[1]) + ) + imgs[line_to_point_assignments[line_idx]][rr, cc] += 1 + return final_points, imgs, points_per_line -def get_top_n_points(line_representations,n,width,threshold=3,mn=4,lines_per_step=2): - final_points=[] - line_to_point_assignments=np.zeros(len(line_representations),dtype=int)-1 - line_to_point_distances=np.zeros(len(line_representations),dtype=int)-1 - for point_idx in range(n): - img=np.zeros((width+1,width+1)) - for line_idx in np.arange(len(line_representations)): - if line_to_point_assignments[line_idx]==-1 or line_to_point_distances[line_idx]>=threshold: - line=line_representations[line_idx] - if len(final_points)>0: # check if its close to the last line - b,m,point,mvec=line_to_mx(line) - fp=final_points[-1] - d=(fp[0]-point[0],fp[1]-point[1]) - if np.sign(mvec[0])==np.sign(d[0]) and np.sign(mvec[1])==np.sign(d[1]): # check if its on the right side - mvec_orthogonal=(-mvec[1],mvec[0]) - distance_to_line=abs(np.dot(d,mvec_orthogonal)/np.linalg.norm(mvec)) - if line_to_point_assignments[line_idx]==-1 or distance_to_line=threshold: - b,m,point,mvec=line_to_mx(line) - bp=boundary_point(b,m,point,width,mvec) - rr,cc=skimage.draw.line( - int(point[0]), int(point[1]), - int(bp[0]),int(bp[1])) - img[rr,cc]+=1 - if img.max()>=mn: - idx=img.argmax() - p=(idx//(width+1),idx%(width+1)) - final_points.append(p) - final_points=np.array(final_points) - points_per_line=np.zeros((len(line_representations)//lines_per_step,2))+width/2 - for output_idx in np.arange(len(line_representations)//lines_per_step): - line_idx=output_idx*lines_per_step - closest_idx=-1 - for _line_idx in range(lines_per_step): - if line_to_point_assignments[line_idx+_line_idx]!=-1 and ( - closest_idx==-1 or line_to_point_distances[line_idx+closest_idx]>line_to_point_distances[line_idx+_line_idx]): - closest_idx=_line_idx - if closest_idx!=-1: - points_per_line[output_idx]=final_points[line_to_point_assignments[line_idx+closest_idx]] - imgs=np.zeros((n,width+1,width+1)) - for point_idx in range(n): - for line_idx in np.arange(len(line_representations)): - if line_to_point_assignments[line_idx]>=0: - line=line_representations[line_idx] - if line_to_point_assignments[line_idx]>=0: - b,m,point,mvec=line_to_mx(line) - bp=boundary_point(b,m,point,width,mvec) - rr,cc=skimage.draw.line( - int(point[0]), int(point[1]), - int(bp[0]),int(bp[1])) - imgs[line_to_point_assignments[line_idx]][rr,cc]+=1 - return final_points,imgs,points_per_line -def boundary_point(b,m,point,width,mvec): - y_xmax=m*width+b - y_xmin=b - x_ymin=-b/m - x_ymax=(width-b)/m - if mvec[0]>0: - if mvec[1]>0: - p1=(x_ymax,width) - p2=(width,y_xmax) - elif mvec[1]<0: - p1=(x_ymin,0) - p2=(width,y_xmax) - else: - p1=(width,point[1]) - p2=(width,point[1]) - elif mvec[0]<0: - if mvec[1]>0: - p1=(0,y_xmin) - p2=(x_ymax,width) - elif mvec[1]<0: - p1=(0,y_xmin) - p2=(x_ymin,0) - else: - p1=(0,point[1]) - p2=(0,point[1]) - else: - if mvec[1]>0: - p1=(point[0],width) - p2=(point[0],width) - elif mvec[1]<0: - p1=(point[0],0) - p2=(point[0],0) - else: - assert(False) - p=p1 - if p1[0]<0 or p1[0]>width or p1[1]<0 or p1[1]>width: - p=p2 - assert(p[0]>=0 and p[0]<=width and p[1]>=0 and p[1]<=width) - return p +def boundary_point(b, m, point, width, mvec): + y_xmax = m * width + b + y_xmin = b + x_ymin = -b / m + x_ymax = (width - b) / m + if mvec[0] > 0: + if mvec[1] > 0: + p1 = (x_ymax, width) + p2 = (width, y_xmax) + elif mvec[1] < 0: + p1 = (x_ymin, 0) + p2 = (width, y_xmax) + else: + p1 = (width, point[1]) + p2 = (width, point[1]) + elif mvec[0] < 0: + if mvec[1] > 0: + p1 = (0, y_xmin) + p2 = (x_ymax, width) + elif mvec[1] < 0: + p1 = (0, y_xmin) + p2 = (x_ymin, 0) + else: + p1 = (0, point[1]) + p2 = (0, point[1]) + else: + if mvec[1] > 0: + p1 = (point[0], width) + p2 = (point[0], width) + elif mvec[1] < 0: + p1 = (point[0], 0) + p2 = (point[0], 0) + else: + assert False + p = p1 + if p1[0] < 0 or p1[0] > width or p1[1] < 0 or p1[1] > width: + p = p2 + assert p[0] >= 0 and p[0] <= width and p[1] >= 0 and p[1] <= width + return p -def lines_to_points(lines,t): - lines=[ line_to_mx(line) for line in lines ] - line_to_points=[] - for line_idx_a in np.arange(len(lines)): - rng = np.random.default_rng(12345+line_idx_a*1337) - a_i,a_m,(x1,y1),_=lines[line_idx_a] - points_for_this_line=[] - for line_idx_b in rng.choice(np.arange(t),size=min(30,t),replace=False): - if line_idx_b>=len(lines): - continue - if line_idx_a==line_idx_b: - continue - b_i,b_m,(x2,y2),_=lines[line_idx_b] - #compute the intercept - if a_i!=b_m and a_m!=b_m: - _x=(b_i-a_i)/(a_m-b_m) - _y=a_m*((b_i-a_i)/(a_m-b_m))+a_i - #check if the point is valid - if a_m<0 and _x0 and _x>x1: - pass - elif b_m<0 and _x0 and _x>x2: - pass - elif ((x1-_x)**2+(y1-_y)**2)<800: - pass - elif ((x2-_x)**2+(y2-_y)**2)<800: - pass - else: - points_for_this_line.append((_x,_y)) - line_to_points.append(points_for_this_line) - return line_to_points +def lines_to_points(lines, t): + lines = [line_to_mx(line) for line in lines] + line_to_points = [] + for line_idx_a in np.arange(len(lines)): + rng = np.random.default_rng(12345 + line_idx_a * 1337) + a_i, a_m, (x1, y1), _ = lines[line_idx_a] + points_for_this_line = [] + for line_idx_b in rng.choice(np.arange(t), size=min(30, t), replace=False): + if line_idx_b >= len(lines): + continue + if line_idx_a == line_idx_b: + continue + b_i, b_m, (x2, y2), _ = lines[line_idx_b] + # compute the intercept + if a_i != b_m and a_m != b_m: + _x = (b_i - a_i) / (a_m - b_m) + _y = a_m * ((b_i - a_i) / (a_m - b_m)) + a_i + # check if the point is valid + if a_m < 0 and _x < x1: + pass + elif a_m > 0 and _x > x1: + pass + elif b_m < 0 and _x < x2: + pass + elif b_m > 0 and _x > x2: + pass + elif ((x1 - _x) ** 2 + (y1 - _y) ** 2) < 800: + pass + elif ((x2 - _x) ** 2 + (y2 - _y) ** 2) < 800: + pass + else: + points_for_this_line.append((_x, _y)) + line_to_points.append(points_for_this_line) + return line_to_points diff --git a/software/model_training_and_inference/utils/image_utils.py b/software/model_training_and_inference/utils/image_utils.py index 63d2363c..368283a3 100644 --- a/software/model_training_and_inference/utils/image_utils.py +++ b/software/model_training_and_inference/utils/image_utils.py @@ -7,77 +7,102 @@ @lru_cache def get_grid(width): - _xy=np.arange(width).astype(np.int16) - _x,_y=np.meshgrid(_xy,_xy) - return np.stack((_x,_y)).transpose(2,1,0) - #return np.concatenate([_y[None],_x[None]]).transpose(0,2) - -#input b,s,2 output: b,s,1,width,width -def detector_positions_to_distance(detector_positions,width): - diffs=get_grid(width)[None,None]-detector_positions[:,:,None,None].astype(np.float32) - return (np.sqrt(np.power(diffs, 2).sum(axis=4)))[:,:,None] # batch, snapshot, 1,x ,y - -#input b,s,2 output: b,s,1,width,width -def detector_positions_to_theta_grid(detector_positions,width,img_width): - bin_size=np.ceil(width/img_width) - diffs=get_grid(img_width)[None,None]*bin_size-detector_positions[:,:,None,None].astype(np.float32) - #return (np.arctan2(diffs[...,1],diffs[...,0]))[:,:,None] # batch, snapshot,1, x ,y - return (np.arctan2(diffs[...,0],diffs[...,1]))[:,:,None] # batch, snapshot,1, x ,y # get angle from x=0, y+ to the right - + _xy = np.arange(width).astype(np.int16) + _x, _y = np.meshgrid(_xy, _xy) + return np.stack((_x, _y)).transpose(2, 1, 0) + # return np.concatenate([_y[None],_x[None]]).transpose(0,2) + + +# input b,s,2 output: b,s,1,width,width +def detector_positions_to_distance(detector_positions, width): + diffs = get_grid(width)[None, None] - detector_positions[:, :, None, None].astype( + np.float32 + ) + return (np.sqrt(np.power(diffs, 2).sum(axis=4)))[ + :, :, None + ] # batch, snapshot, 1,x ,y + + +# input b,s,2 output: b,s,1,width,width +def detector_positions_to_theta_grid(detector_positions, width, img_width): + bin_size = np.ceil(width / img_width) + diffs = get_grid(img_width)[None, None] * bin_size - detector_positions[ + :, :, None, None + ].astype(np.float32) + # return (np.arctan2(diffs[...,1],diffs[...,0]))[:,:,None] # batch, snapshot,1, x ,y + return (np.arctan2(diffs[..., 0], diffs[..., 1]))[ + :, :, None + ] # batch, snapshot,1, x ,y # get angle from x=0, y+ to the right + + def blur2(img): - blur=torchvision.transforms.GaussianBlur(11, sigma=8.0) - return blur(blur(img)) + blur = torchvision.transforms.GaussianBlur(11, sigma=8.0) + return blur(blur(img)) + def blur5(img): - blur=torchvision.transforms.GaussianBlur(11, sigma=8.0) - return blur(blur(blur(blur(blur(img))))) + blur = torchvision.transforms.GaussianBlur(11, sigma=8.0) + return blur(blur(blur(blur(blur(img))))) + + def blur10(img): - return blur5(blur5(img)) - -def labels_to_source_images(labels,width,img_width=128): - b,s,n_sources,_=labels.shape - offset=0 #50 # takes too much compute! - label_images=torch.zeros(( - b,s, - img_width+2*offset, - img_width+2*offset)) - bin_size=np.ceil(width/img_width) - for b_idx in np.arange(b): - for s_idx in np.arange(s): - for source_idx in np.arange(n_sources): - source_x,source_y=labels[b_idx,s_idx,source_idx] - if s_idx==0: - print(labels[b_idx,s_idx,source_idx]) - print(int(source_x/bin_size)+offset,int(source_y/bin_size)+offset) - label_images[b_idx,s_idx,int(source_x/bin_size)+offset,int(source_y/bin_size)+offset]=1 - - label_images=blur5( - label_images.reshape( - b*s,1, - img_width+2*offset, - img_width+2*offset)).reshape( - b,s,1, - img_width+2*offset, - img_width+2*offset) - #label_images=label_images[...,offset:-offset,offset:-offset] # trim the rest - assert(label_images.shape[3]==img_width) - label_images=label_images/label_images.sum(axis=[3,4],keepdims=True) - return label_images - -def radio_to_image(beam_former_outputs_at_t,theta_at_pos,detector_orientation): - #theta_at_pos=(theta_at_pos+np.pi-detector_orientation[...,None,None])%(2*np.pi) - #theta_at_pos=(theta_at_pos+detector_orientation[...,None,None])%(2*np.pi) - theta_at_pos=(theta_at_pos-detector_orientation[...,None,None])%(2*np.pi) - #theta_idxs=(((theta_at_pos+np.pi)/(2*np.pi))*(beam_former_outputs_at_t.shape[-1]-1)).round().astype(int) - theta_idxs=(((theta_at_pos/(2*np.pi)+0.5)%1.0)*(beam_former_outputs_at_t.shape[-1]-1)).round().astype(int) - b,s,_,width,_=theta_at_pos.shape - #return theta_at_pos.reshape(b,s,1,width,width) - beam_former_outputs_at_t=beam_former_outputs_at_t#[:]/beam_former_outputs_at_t.sum(axis=2,keepdims=True) - outputs=[] - for b_idx in np.arange(b): - for s_idx in np.arange(s): - outputs.append( - np.take( - beam_former_outputs_at_t[b_idx,s_idx], - theta_idxs[b_idx,s_idx].reshape(-1)).reshape(theta_idxs[b_idx,s_idx].shape)) - return np.stack(outputs).reshape(b,s,1,width,width) + return blur5(blur5(img)) + + +def labels_to_source_images(labels, width, img_width=128): + b, s, n_sources, _ = labels.shape + offset = 0 # 50 # takes too much compute! + label_images = torch.zeros((b, s, img_width + 2 * offset, img_width + 2 * offset)) + bin_size = np.ceil(width / img_width) + for b_idx in np.arange(b): + for s_idx in np.arange(s): + for source_idx in np.arange(n_sources): + source_x, source_y = labels[b_idx, s_idx, source_idx] + if s_idx == 0: + print(labels[b_idx, s_idx, source_idx]) + print( + int(source_x / bin_size) + offset, + int(source_y / bin_size) + offset, + ) + label_images[ + b_idx, + s_idx, + int(source_x / bin_size) + offset, + int(source_y / bin_size) + offset, + ] = 1 + + label_images = blur5( + label_images.reshape(b * s, 1, img_width + 2 * offset, img_width + 2 * offset) + ).reshape(b, s, 1, img_width + 2 * offset, img_width + 2 * offset) + # label_images=label_images[...,offset:-offset,offset:-offset] # trim the rest + assert label_images.shape[3] == img_width + label_images = label_images / label_images.sum(axis=[3, 4], keepdims=True) + return label_images + + +def radio_to_image(beam_former_outputs_at_t, theta_at_pos, detector_orientation): + # theta_at_pos=(theta_at_pos+np.pi-detector_orientation[...,None,None])%(2*np.pi) + # theta_at_pos=(theta_at_pos+detector_orientation[...,None,None])%(2*np.pi) + theta_at_pos = (theta_at_pos - detector_orientation[..., None, None]) % (2 * np.pi) + # theta_idxs=(((theta_at_pos+np.pi)/(2*np.pi))*(beam_former_outputs_at_t.shape[-1]-1)).round().astype(int) + theta_idxs = ( + ( + ((theta_at_pos / (2 * np.pi) + 0.5) % 1.0) + * (beam_former_outputs_at_t.shape[-1] - 1) + ) + .round() + .astype(int) + ) + b, s, _, width, _ = theta_at_pos.shape + # return theta_at_pos.reshape(b,s,1,width,width) + beam_former_outputs_at_t = beam_former_outputs_at_t # [:]/beam_former_outputs_at_t.sum(axis=2,keepdims=True) + outputs = [] + for b_idx in np.arange(b): + for s_idx in np.arange(s): + outputs.append( + np.take( + beam_former_outputs_at_t[b_idx, s_idx], + theta_idxs[b_idx, s_idx].reshape(-1), + ).reshape(theta_idxs[b_idx, s_idx].shape) + ) + return np.stack(outputs).reshape(b, s, 1, width, width) diff --git a/software/model_training_and_inference/utils/plot.py b/software/model_training_and_inference/utils/plot.py index 2d85b4a5..1887e337 100644 --- a/software/model_training_and_inference/utils/plot.py +++ b/software/model_training_and_inference/utils/plot.py @@ -3,310 +3,492 @@ import torch from PIL import Image import skimage -from utils.image_utils import (detector_positions_to_theta_grid, - labels_to_source_images, radio_to_image) +from utils.image_utils import ( + detector_positions_to_theta_grid, + labels_to_source_images, + radio_to_image, +) from utils.baseline_algorithm import get_top_n_peaks, baseline_algorithm -#depends if we are using y=0 or x=0 as origin -#if x=0 (y+) then we want x=sin, y=cos -#if y=0 (x+) then we want x=cos, y=sin +# depends if we are using y=0 or x=0 as origin +# if x=0 (y+) then we want x=sin, y=cos +# if y=0 (x+) then we want x=cos, y=sin def get_xy_from_theta(theta): - return np.array([ - np.sin(theta), - np.cos(theta), - ]).T - -def plot_predictions_and_baseline(session,args,step,pred_a,pred_b): - width=session['width_at_t'][0][0] - - fig=plt.figure(figsize=(5*2,5*2)) - axs=fig.subplots(2,2) - - plot_trajectory(axs[0,0],session['detector_position_at_t'][:step],width,ms=30,label='detector') - plot_trajectory(axs[0,1],session['detector_position_at_t'][:step],width,ms=30,label='detector') - - #plot directions on the the space diagram - direction=session['detector_position_at_t'][step]+0.25*session['width_at_t'][0]*get_xy_from_theta(session['detector_orientation_at_t'][step]) - axs[0,0].plot( - [session['detector_position_at_t'][step][0],direction[0,0]], - [session['detector_position_at_t'][step][1],direction[0,1]]) - anti_direction=session['detector_position_at_t'][step]+0.25*session['width_at_t'][0]*get_xy_from_theta(session['detector_orientation_at_t'][step]+np.pi/2) - axs[0,0].plot( - [session['detector_position_at_t'][step][0],anti_direction[0,0]], - [session['detector_position_at_t'][step][1],anti_direction[0,1]]) - - #for every time step plot the lines - lines=[] - for idx in range(step+1): - _thetas=session['thetas_at_t'][idx][get_top_n_peaks(session['beam_former_outputs_at_t'][idx])] - for _theta in _thetas: - direction=get_xy_from_theta(session['detector_orientation_at_t'][idx]+_theta) - emitter_direction=session['detector_position_at_t'][idx]+2.0*session['width_at_t'][0]*direction - lines.append(([session['detector_position_at_t'][idx][0],emitter_direction[0,0]], - [session['detector_position_at_t'][idx][1],emitter_direction[0,1]])) - - for x,y in lines: - axs[0,1].plot(x,y,c='blue',linewidth=4,alpha=0.1) - - for n in np.arange(session['source_positions_at_t'].shape[1]): - #rings=(session['broadcasting_positions_at_t'][idx,n,0]==1) - rings=False - plot_trajectory(axs[0,0],session['source_positions_at_t'][:idx,n],width,ms=15,c='r',rings=rings,label='emitter %d' % n) - axs[0,0].set_title("Position map") - axs[0,1].set_title("Direction estimates") - for x in [0,1]: - handles, labels = axs[0,x].get_legend_handles_labels() - by_label = dict(zip(labels, handles)) - axs[0,x].legend(by_label.values(), by_label.keys()) - for y in [0,1]: - #axs[y,x].set_title("Position map") - axs[y,x].set_xlabel("X (m)") - axs[y,x].set_ylabel("Y (m)") - axs[y,x].set_xlim([0,width]) - axs[y,x].set_ylim([0,width]) - - true_positions=session['source_positions_at_t'][session['broadcasting_positions_at_t'].astype(bool)[...,0]] - true_positions_noise=true_positions+np.random.randn(*true_positions.shape)*3 - - for ax_idx,pred in [(0,pred_a),(1,pred_b)]: - axs[1,ax_idx].set_title("error in %s" % pred['name']) - axs[1,ax_idx].scatter(pred['predictions'][:,0],pred['predictions'][:,1],s=15,c='r',alpha=0.1) - for idx in np.arange(step): - _x,_y=pred['predictions'][idx]+np.random.randn(2) - x,y=true_positions_noise[idx] - axs[1,ax_idx].plot([_x,x],[_y,y],color='black',linewidth=1,alpha=0.1) - - - fn='%s_%04d_lines.png' % (args.output_prefix,step) - fig.savefig(fn) - plt.close(fig) - return fn - -def plot_space(ax,session): - width=session['width_at_t'][0] - ax.set_xlim([0,width]) - ax.set_ylim([0,width]) - - markers=['o','v','D'] - colors=['g', 'b', 'y'] - for receiver_idx in np.arange(session['receiver_positions_at_t'].shape[1]): - ax.scatter(session['receiver_positions_at_t'][:,receiver_idx,0],session['receiver_positions_at_t'][:,receiver_idx,1],label="Receiver %d" % receiver_idx ,facecolors='none',marker=markers[receiver_idx%len(markers)],edgecolor=colors[receiver_idx%len(colors)]) - for source_idx in np.arange(session['source_positions_at_t'].shape[1]): - ax.scatter(session['source_positions_at_t'][:,source_idx,0],session['source_positions_at_t'][:,source_idx,1],label="Source %d" % source_idx ,facecolors='none',marker=markers[source_idx%len(markers)],edgecolor='r') - ax.legend() - ax.set_xlabel("x (m)") - ax.set_ylabel("y (m)") - -def plot_trajectory(ax, - positions, - width, - ms=30, - steps_per_fade=10, - fadep=0.8, - c='b', - rings=False, - label=None): - ax.set_xlim([0,width]) - ax.set_ylim([0,width]) - n_steps=positions.shape[0]//steps_per_fade - if positions.shape[0]%steps_per_fade!=0: - n_steps+=1 - - alpha=fadep - for n in np.arange(1,n_steps): - start=positions.shape[0]-(n+1)*steps_per_fade - end=start+steps_per_fade - start=max(0,start) - ax.plot( positions[start:end,0], positions[start:end,1],'--',alpha=alpha,color=c,label=label) - alpha*=fadep - start=positions.shape[0]-steps_per_fade - end=start+steps_per_fade - ax.plot( positions[start:end,0], positions[start:end,1],'--',alpha=1.0,color=c,label=label) - - ax.plot( positions[-1,0], positions[-1,1],'.',ms=ms,c=c) - if rings: - n=4 - for x in range(n): - ax.plot( positions[-1,0], positions[-1,1],'.',ms=ms*(1.8**x),c=c,alpha=1/n) - - -def plot_lines(session,steps,output_prefix): - width=session['width_at_t'][0][0] - - #extract the images - d={} - filenames=[] - plt.ioff() - lines=[] - for idx in np.arange(1,steps): - fig=plt.figure(figsize=(9*3,9)) - axs=fig.subplots(1,3) - - plot_trajectory(axs[0],session['detector_position_at_t'][:idx],width,ms=30,label='detector') - plot_trajectory(axs[1],session['detector_position_at_t'][:idx],width,ms=30,label='detector') - direction=session['detector_position_at_t'][idx]+0.25*session['width_at_t'][0]*get_xy_from_theta(session['detector_orientation_at_t'][idx]) - axs[0].plot( - [session['detector_position_at_t'][idx][0],direction[0,0]], - [session['detector_position_at_t'][idx][1],direction[0,1]]) - anti_direction=session['detector_position_at_t'][idx]+0.25*session['width_at_t'][0]*get_xy_from_theta(session['detector_orientation_at_t'][idx]+np.pi/2) - axs[0].plot( - [session['detector_position_at_t'][idx][0],anti_direction[0,0]], - [session['detector_position_at_t'][idx][1],anti_direction[0,1]]) - _thetas=session['thetas_at_t'][idx][get_top_n_peaks(session['beam_former_outputs_at_t'][idx])] - for _theta in _thetas: - direction=get_xy_from_theta(session['detector_orientation_at_t'][idx]+_theta) - emitter_direction=session['detector_position_at_t'][idx]+2.0*session['width_at_t'][0]*direction - lines.append(([session['detector_position_at_t'][idx][0],emitter_direction[0,0]], - [session['detector_position_at_t'][idx][1],emitter_direction[0,1]])) - - for x,y in lines: - axs[0].plot(x,y,c='blue',linewidth=4,alpha=0.1) - - emitter_direction=session['detector_position_at_t'][idx]+0.25*session['width_at_t'][0]*get_xy_from_theta(session['detector_orientation_at_t'][idx]+session['source_theta_at_t'][idx,0]) - axs[0].plot( - [session['detector_position_at_t'][idx][0],emitter_direction[0,0]], - [session['detector_position_at_t'][idx][1],emitter_direction[0,1]]) - - for n in np.arange(session['source_positions_at_t'].shape[1]): - rings=(session['broadcasting_positions_at_t'][idx,n,0]==1) - plot_trajectory(axs[0],session['source_positions_at_t'][:idx,n],width,ms=15,c='r',rings=rings,label='emitter %d' % n) - handles, labels = axs[0].get_legend_handles_labels() - by_label = dict(zip(labels, handles)) - axs[0].legend(by_label.values(), by_label.keys()) - - axs[2].plot(session['thetas_at_t'][idx],session['beam_former_outputs_at_t'][idx]) - axs[2].axvline(x=session['source_theta_at_t'][idx,0],c='r') - axs[2].set_title("Beamformer output at t=%d" % idx) - axs[2].set_xlabel("Theta (rel. to detector)") - axs[2].set_ylabel("Signal strength") - axs[0].set_title("Position map") - axs[0].set_xlabel("X (m)") - axs[0].set_ylabel("Y (m)") - axs[1].set_title("Guess map") - axs[1].set_xlabel("X (m)") - axs[1].set_ylabel("Y (m)") - - fp,imgs,pred_points=baseline_algorithm(session,width,steps=idx) - for pred_point in pred_points: - axs[0].scatter([pred_point[0]],[pred_point[1]]) - axs[1].imshow(imgs[:3].transpose([2,1,0])/imgs.max()) - colors=['r','green','blue'] - for _idx in range(min(3,len(fp))): - axs[1].scatter([fp[_idx][0]],[fp[_idx][1]],color=colors[_idx],s=900) - - - fn='%s_%04d_lines.png' % (output_prefix,idx) - filenames.append(fn) + return np.array( + [ + np.sin(theta), + np.cos(theta), + ] + ).T + + +def plot_predictions_and_baseline(session, args, step, pred_a, pred_b): + width = session["width_at_t"][0][0] + + fig = plt.figure(figsize=(5 * 2, 5 * 2)) + axs = fig.subplots(2, 2) + + plot_trajectory( + axs[0, 0], + session["detector_position_at_t"][:step], + width, + ms=30, + label="detector", + ) + plot_trajectory( + axs[0, 1], + session["detector_position_at_t"][:step], + width, + ms=30, + label="detector", + ) + + # plot directions on the the space diagram + direction = session["detector_position_at_t"][step] + 0.25 * session["width_at_t"][ + 0 + ] * get_xy_from_theta(session["detector_orientation_at_t"][step]) + axs[0, 0].plot( + [session["detector_position_at_t"][step][0], direction[0, 0]], + [session["detector_position_at_t"][step][1], direction[0, 1]], + ) + anti_direction = session["detector_position_at_t"][step] + 0.25 * session[ + "width_at_t" + ][0] * get_xy_from_theta(session["detector_orientation_at_t"][step] + np.pi / 2) + axs[0, 0].plot( + [session["detector_position_at_t"][step][0], anti_direction[0, 0]], + [session["detector_position_at_t"][step][1], anti_direction[0, 1]], + ) + + # for every time step plot the lines + lines = [] + for idx in range(step + 1): + _thetas = session["thetas_at_t"][idx][ + get_top_n_peaks(session["beam_former_outputs_at_t"][idx]) + ] + for _theta in _thetas: + direction = get_xy_from_theta( + session["detector_orientation_at_t"][idx] + _theta + ) + emitter_direction = ( + session["detector_position_at_t"][idx] + + 2.0 * session["width_at_t"][0] * direction + ) + lines.append( + ( + [ + session["detector_position_at_t"][idx][0], + emitter_direction[0, 0], + ], + [ + session["detector_position_at_t"][idx][1], + emitter_direction[0, 1], + ], + ) + ) + + for x, y in lines: + axs[0, 1].plot(x, y, c="blue", linewidth=4, alpha=0.1) + + for n in np.arange(session["source_positions_at_t"].shape[1]): + # rings=(session['broadcasting_positions_at_t'][idx,n,0]==1) + rings = False + plot_trajectory( + axs[0, 0], + session["source_positions_at_t"][:idx, n], + width, + ms=15, + c="r", + rings=rings, + label="emitter %d" % n, + ) + axs[0, 0].set_title("Position map") + axs[0, 1].set_title("Direction estimates") + for x in [0, 1]: + handles, labels = axs[0, x].get_legend_handles_labels() + by_label = dict(zip(labels, handles)) + axs[0, x].legend(by_label.values(), by_label.keys()) + for y in [0, 1]: + # axs[y,x].set_title("Position map") + axs[y, x].set_xlabel("X (m)") + axs[y, x].set_ylabel("Y (m)") + axs[y, x].set_xlim([0, width]) + axs[y, x].set_ylim([0, width]) + + true_positions = session["source_positions_at_t"][ + session["broadcasting_positions_at_t"].astype(bool)[..., 0] + ] + true_positions_noise = true_positions + np.random.randn(*true_positions.shape) * 3 + + for ax_idx, pred in [(0, pred_a), (1, pred_b)]: + axs[1, ax_idx].set_title("error in %s" % pred["name"]) + axs[1, ax_idx].scatter( + pred["predictions"][:, 0], pred["predictions"][:, 1], s=15, c="r", alpha=0.1 + ) + for idx in np.arange(step): + _x, _y = pred["predictions"][idx] + np.random.randn(2) + x, y = true_positions_noise[idx] + axs[1, ax_idx].plot([_x, x], [_y, y], color="black", linewidth=1, alpha=0.1) + + fn = "%s_%04d_lines.png" % (args.output_prefix, step) fig.savefig(fn) plt.close(fig) - plt.ion() - return filenames - -#generate the images for the session -def plot_full_session(session,steps,output_prefix,img_width=128,invert=False): - width=session['width_at_t'][0][0] - - #extract the images - d={} - d['source_image_at_t']=labels_to_source_images(torch.from_numpy(session['source_positions_at_t'])[None],width,img_width=img_width)[0] - d['detector_theta_image_at_t']=detector_positions_to_theta_grid( - session['detector_position_at_t'][None],width,img_width=img_width)[0] - d['radio_image_at_t']=radio_to_image( - session['beam_former_outputs_at_t'][None], - d['detector_theta_image_at_t'][None], - session['detector_orientation_at_t'][None])[0] - d['radio_image_at_t_normed']=d['radio_image_at_t']/d['radio_image_at_t'].sum(axis=2,keepdims=True).sum(axis=3,keepdims=True) - filenames=[] - plt.ioff() - for idx in np.arange(1,steps): - fig=plt.figure(figsize=(12,12)) - axs=fig.subplots(2,2) - for _a in [0,1]: - for _b in [0,1]: - if _a==0 and _b==1: - continue - axs[_a,_b].set_xlabel("X (m)") - axs[_a,_b].set_ylabel("Y (m)") - - axs[0,0].set_title("Position map") - plot_trajectory(axs[0,0],session['detector_position_at_t'][:idx+1],width,ms=30,label='detector') - direction=session['detector_position_at_t'][idx]+0.25*session['width_at_t'][0]*get_xy_from_theta(session['detector_orientation_at_t'][idx]) - - axs[0,0].plot( - [session['detector_position_at_t'][idx][0],direction[0,0]], - [session['detector_position_at_t'][idx][1],direction[0,1]]) - - anti_direction=session['detector_position_at_t'][idx]+0.25*session['width_at_t'][0]*get_xy_from_theta(session['detector_orientation_at_t'][idx]+np.pi/2) - - axs[0,0].plot( - [session['detector_position_at_t'][idx][0],anti_direction[0,0]], - [session['detector_position_at_t'][idx][1],anti_direction[0,1]]) - - emitter_direction=session['detector_position_at_t'][idx]+0.25*session['width_at_t'][0]*get_xy_from_theta(session['detector_orientation_at_t'][idx]+session['source_theta_at_t'][idx,0]) - axs[0,0].plot( - [session['detector_position_at_t'][idx][0],emitter_direction[0,0]], - [session['detector_position_at_t'][idx][1],emitter_direction[0,1]]) - for n in np.arange(session['source_positions_at_t'].shape[1]): - rings=(session['broadcasting_positions_at_t'][idx,n,0]==1) - plot_trajectory(axs[0,0],session['source_positions_at_t'][:idx+1,n],width,ms=15,c='r',rings=rings,label='emitter %d' % n) - - #axs[0,0].legend() - handles, labels = axs[0,0].get_legend_handles_labels() - by_label = dict(zip(labels, handles)) - axs[0,0].legend(by_label.values(), by_label.keys()) - - #lets draw the radio - #axs[1,0].imshow(d['source_image_at_t'][idx,0].T, - # origin='upper', - # extent=(0, - # d['source_image_at_t'][idx,0].shape[0], - # d['source_image_at_t'][idx,0].shape[1], - # 0) - #) #,origin='lower') - axs[1,0].imshow(d['source_image_at_t'][idx,0].T) - axs[1,0].set_title("Emitters as image at t=%d" % idx) - - axs[1,1].imshow( - d['radio_image_at_t'][idx,0].T) - if invert: - axs[0,0].invert_xaxis() - axs[0,0].invert_yaxis() - axs[1,0].invert_xaxis() - axs[1,1].invert_xaxis() - else: - #axs[1,0].invert_yaxis() - axs[1,1].invert_yaxis() - #axs[1,1].invert_xaxis() - pass - # origin='upper', - # extent=(0, - # d['radio_image_at_t'][idx,0].shape[0], - # d['radio_image_at_t'][idx,0].shape[1], - # 0)) - #d['radio_image_at_t']=radio_to_image(session['beam_former_outputs_at_t'][None],d['detector_theta_image_at_t'][None],session['detector_orientation_at_t'][None])[0] - #axs[1,1].imshow(d['detector_theta_image_at_t'][0,0]) - axs[1,1].set_title("Radio feature at t=%d" % idx) - axs[0,1].plot(session['thetas_at_t'][idx],session['beam_former_outputs_at_t'][idx]) - axs[0,1].axvline(x=session['source_theta_at_t'][idx,0],c='r') - axs[0,1].set_title("Beamformer output at t=%d" % idx) - axs[0,1].set_xlabel("Theta (rel. to detector)") - axs[0,1].set_ylabel("Signal strength") - - fn='%s_%04d.png' % (output_prefix,idx) - filenames.append(fn) - fig.savefig(fn) - plt.close(fig) - plt.ion() - return filenames - - -def filenames_to_gif(filenames,output_gif_fn,size=(600,600),duration=200): - images=[] - for fn in filenames: - images.append(Image.open(fn).resize(size)) - - images[0].save(output_gif_fn, - save_all = True, append_images = images[1:], - optimize = False, duration = duration,loop=0) + return fn + + +def plot_space(ax, session): + width = session["width_at_t"][0] + ax.set_xlim([0, width]) + ax.set_ylim([0, width]) + + markers = ["o", "v", "D"] + colors = ["g", "b", "y"] + for receiver_idx in np.arange(session["receiver_positions_at_t"].shape[1]): + ax.scatter( + session["receiver_positions_at_t"][:, receiver_idx, 0], + session["receiver_positions_at_t"][:, receiver_idx, 1], + label="Receiver %d" % receiver_idx, + facecolors="none", + marker=markers[receiver_idx % len(markers)], + edgecolor=colors[receiver_idx % len(colors)], + ) + for source_idx in np.arange(session["source_positions_at_t"].shape[1]): + ax.scatter( + session["source_positions_at_t"][:, source_idx, 0], + session["source_positions_at_t"][:, source_idx, 1], + label="Source %d" % source_idx, + facecolors="none", + marker=markers[source_idx % len(markers)], + edgecolor="r", + ) + ax.legend() + ax.set_xlabel("x (m)") + ax.set_ylabel("y (m)") + + +def plot_trajectory( + ax, + positions, + width, + ms=30, + steps_per_fade=10, + fadep=0.8, + c="b", + rings=False, + label=None, +): + ax.set_xlim([0, width]) + ax.set_ylim([0, width]) + n_steps = positions.shape[0] // steps_per_fade + if positions.shape[0] % steps_per_fade != 0: + n_steps += 1 + + alpha = fadep + for n in np.arange(1, n_steps): + start = positions.shape[0] - (n + 1) * steps_per_fade + end = start + steps_per_fade + start = max(0, start) + ax.plot( + positions[start:end, 0], + positions[start:end, 1], + "--", + alpha=alpha, + color=c, + label=label, + ) + alpha *= fadep + start = positions.shape[0] - steps_per_fade + end = start + steps_per_fade + ax.plot( + positions[start:end, 0], + positions[start:end, 1], + "--", + alpha=1.0, + color=c, + label=label, + ) + + ax.plot(positions[-1, 0], positions[-1, 1], ".", ms=ms, c=c) + if rings: + n = 4 + for x in range(n): + ax.plot( + positions[-1, 0], + positions[-1, 1], + ".", + ms=ms * (1.8**x), + c=c, + alpha=1 / n, + ) + + +def plot_lines(session, steps, output_prefix): + width = session["width_at_t"][0][0] + + # extract the images + d = {} + filenames = [] + plt.ioff() + lines = [] + for idx in np.arange(1, steps): + fig = plt.figure(figsize=(9 * 3, 9)) + axs = fig.subplots(1, 3) + + plot_trajectory( + axs[0], + session["detector_position_at_t"][:idx], + width, + ms=30, + label="detector", + ) + plot_trajectory( + axs[1], + session["detector_position_at_t"][:idx], + width, + ms=30, + label="detector", + ) + direction = session["detector_position_at_t"][idx] + 0.25 * session[ + "width_at_t" + ][0] * get_xy_from_theta(session["detector_orientation_at_t"][idx]) + axs[0].plot( + [session["detector_position_at_t"][idx][0], direction[0, 0]], + [session["detector_position_at_t"][idx][1], direction[0, 1]], + ) + anti_direction = session["detector_position_at_t"][idx] + 0.25 * session[ + "width_at_t" + ][0] * get_xy_from_theta(session["detector_orientation_at_t"][idx] + np.pi / 2) + axs[0].plot( + [session["detector_position_at_t"][idx][0], anti_direction[0, 0]], + [session["detector_position_at_t"][idx][1], anti_direction[0, 1]], + ) + _thetas = session["thetas_at_t"][idx][ + get_top_n_peaks(session["beam_former_outputs_at_t"][idx]) + ] + for _theta in _thetas: + direction = get_xy_from_theta( + session["detector_orientation_at_t"][idx] + _theta + ) + emitter_direction = ( + session["detector_position_at_t"][idx] + + 2.0 * session["width_at_t"][0] * direction + ) + lines.append( + ( + [ + session["detector_position_at_t"][idx][0], + emitter_direction[0, 0], + ], + [ + session["detector_position_at_t"][idx][1], + emitter_direction[0, 1], + ], + ) + ) + + for x, y in lines: + axs[0].plot(x, y, c="blue", linewidth=4, alpha=0.1) + + emitter_direction = session["detector_position_at_t"][idx] + 0.25 * session[ + "width_at_t" + ][0] * get_xy_from_theta( + session["detector_orientation_at_t"][idx] + + session["source_theta_at_t"][idx, 0] + ) + axs[0].plot( + [session["detector_position_at_t"][idx][0], emitter_direction[0, 0]], + [session["detector_position_at_t"][idx][1], emitter_direction[0, 1]], + ) + + for n in np.arange(session["source_positions_at_t"].shape[1]): + rings = session["broadcasting_positions_at_t"][idx, n, 0] == 1 + plot_trajectory( + axs[0], + session["source_positions_at_t"][:idx, n], + width, + ms=15, + c="r", + rings=rings, + label="emitter %d" % n, + ) + handles, labels = axs[0].get_legend_handles_labels() + by_label = dict(zip(labels, handles)) + axs[0].legend(by_label.values(), by_label.keys()) + + axs[2].plot( + session["thetas_at_t"][idx], session["beam_former_outputs_at_t"][idx] + ) + axs[2].axvline(x=session["source_theta_at_t"][idx, 0], c="r") + axs[2].set_title("Beamformer output at t=%d" % idx) + axs[2].set_xlabel("Theta (rel. to detector)") + axs[2].set_ylabel("Signal strength") + axs[0].set_title("Position map") + axs[0].set_xlabel("X (m)") + axs[0].set_ylabel("Y (m)") + axs[1].set_title("Guess map") + axs[1].set_xlabel("X (m)") + axs[1].set_ylabel("Y (m)") + + fp, imgs, pred_points = baseline_algorithm(session, width, steps=idx) + for pred_point in pred_points: + axs[0].scatter([pred_point[0]], [pred_point[1]]) + axs[1].imshow(imgs[:3].transpose([2, 1, 0]) / imgs.max()) + colors = ["r", "green", "blue"] + for _idx in range(min(3, len(fp))): + axs[1].scatter([fp[_idx][0]], [fp[_idx][1]], color=colors[_idx], s=900) + + fn = "%s_%04d_lines.png" % (output_prefix, idx) + filenames.append(fn) + fig.savefig(fn) + plt.close(fig) + plt.ion() + return filenames + + +# generate the images for the session +def plot_full_session(session, steps, output_prefix, img_width=128, invert=False): + width = session["width_at_t"][0][0] + + # extract the images + d = {} + d["source_image_at_t"] = labels_to_source_images( + torch.from_numpy(session["source_positions_at_t"])[None], + width, + img_width=img_width, + )[0] + d["detector_theta_image_at_t"] = detector_positions_to_theta_grid( + session["detector_position_at_t"][None], width, img_width=img_width + )[0] + d["radio_image_at_t"] = radio_to_image( + session["beam_former_outputs_at_t"][None], + d["detector_theta_image_at_t"][None], + session["detector_orientation_at_t"][None], + )[0] + d["radio_image_at_t_normed"] = d["radio_image_at_t"] / d["radio_image_at_t"].sum( + axis=2, keepdims=True + ).sum(axis=3, keepdims=True) + filenames = [] + plt.ioff() + for idx in np.arange(1, steps): + fig = plt.figure(figsize=(12, 12)) + axs = fig.subplots(2, 2) + for _a in [0, 1]: + for _b in [0, 1]: + if _a == 0 and _b == 1: + continue + axs[_a, _b].set_xlabel("X (m)") + axs[_a, _b].set_ylabel("Y (m)") + + axs[0, 0].set_title("Position map") + plot_trajectory( + axs[0, 0], + session["detector_position_at_t"][: idx + 1], + width, + ms=30, + label="detector", + ) + direction = session["detector_position_at_t"][idx] + 0.25 * session[ + "width_at_t" + ][0] * get_xy_from_theta(session["detector_orientation_at_t"][idx]) + + axs[0, 0].plot( + [session["detector_position_at_t"][idx][0], direction[0, 0]], + [session["detector_position_at_t"][idx][1], direction[0, 1]], + ) + + anti_direction = session["detector_position_at_t"][idx] + 0.25 * session[ + "width_at_t" + ][0] * get_xy_from_theta(session["detector_orientation_at_t"][idx] + np.pi / 2) + + axs[0, 0].plot( + [session["detector_position_at_t"][idx][0], anti_direction[0, 0]], + [session["detector_position_at_t"][idx][1], anti_direction[0, 1]], + ) + + emitter_direction = session["detector_position_at_t"][idx] + 0.25 * session[ + "width_at_t" + ][0] * get_xy_from_theta( + session["detector_orientation_at_t"][idx] + + session["source_theta_at_t"][idx, 0] + ) + axs[0, 0].plot( + [session["detector_position_at_t"][idx][0], emitter_direction[0, 0]], + [session["detector_position_at_t"][idx][1], emitter_direction[0, 1]], + ) + for n in np.arange(session["source_positions_at_t"].shape[1]): + rings = session["broadcasting_positions_at_t"][idx, n, 0] == 1 + plot_trajectory( + axs[0, 0], + session["source_positions_at_t"][: idx + 1, n], + width, + ms=15, + c="r", + rings=rings, + label="emitter %d" % n, + ) + + # axs[0,0].legend() + handles, labels = axs[0, 0].get_legend_handles_labels() + by_label = dict(zip(labels, handles)) + axs[0, 0].legend(by_label.values(), by_label.keys()) + + # lets draw the radio + # axs[1,0].imshow(d['source_image_at_t'][idx,0].T, + # origin='upper', + # extent=(0, + # d['source_image_at_t'][idx,0].shape[0], + # d['source_image_at_t'][idx,0].shape[1], + # 0) + # ) #,origin='lower') + axs[1, 0].imshow(d["source_image_at_t"][idx, 0].T) + axs[1, 0].set_title("Emitters as image at t=%d" % idx) + + axs[1, 1].imshow(d["radio_image_at_t"][idx, 0].T) + if invert: + axs[0, 0].invert_xaxis() + axs[0, 0].invert_yaxis() + axs[1, 0].invert_xaxis() + axs[1, 1].invert_xaxis() + else: + # axs[1,0].invert_yaxis() + axs[1, 1].invert_yaxis() + # axs[1,1].invert_xaxis() + pass + # origin='upper', + # extent=(0, + # d['radio_image_at_t'][idx,0].shape[0], + # d['radio_image_at_t'][idx,0].shape[1], + # 0)) + # d['radio_image_at_t']=radio_to_image(session['beam_former_outputs_at_t'][None],d['detector_theta_image_at_t'][None],session['detector_orientation_at_t'][None])[0] + # axs[1,1].imshow(d['detector_theta_image_at_t'][0,0]) + axs[1, 1].set_title("Radio feature at t=%d" % idx) + axs[0, 1].plot( + session["thetas_at_t"][idx], session["beam_former_outputs_at_t"][idx] + ) + axs[0, 1].axvline(x=session["source_theta_at_t"][idx, 0], c="r") + axs[0, 1].set_title("Beamformer output at t=%d" % idx) + axs[0, 1].set_xlabel("Theta (rel. to detector)") + axs[0, 1].set_ylabel("Signal strength") + + fn = "%s_%04d.png" % (output_prefix, idx) + filenames.append(fn) + fig.savefig(fn) + plt.close(fig) + plt.ion() + return filenames + + +def filenames_to_gif(filenames, output_gif_fn, size=(600, 600), duration=200): + images = [] + for fn in filenames: + images.append(Image.open(fn).resize(size)) + + images[0].save( + output_gif_fn, + save_all=True, + append_images=images[1:], + optimize=False, + duration=duration, + loop=0, + ) diff --git a/software/model_training_and_inference/utils/rf.py b/software/model_training_and_inference/utils/rf.py index 3a5ac9a5..1fd960ff 100644 --- a/software/model_training_and_inference/utils/rf.py +++ b/software/model_training_and_inference/utils/rf.py @@ -1,218 +1,308 @@ import matplotlib.pyplot as plt import numpy as np -#from numba import jit + +# from numba import jit import functools -numba=False -''' +numba = False + +""" Given some guess of the source of direction we can shift the carrier frequency phase of received samples at the N different receivers. If the guess of the source direction is correct, the signal from the N different receivers should interfer constructively. -''' +""" + + @functools.lru_cache(maxsize=1024) -def rf_linspace(s,e,i): - return np.linspace(s,e,i) +def rf_linspace(s, e, i): + return np.linspace(s, e, i) -''' + +""" Rotate by orientation If we are left multiplying then its a right (clockwise) rotation -''' +""" + + @functools.lru_cache(maxsize=1024) -def rotation_matrix(orientation): - s = np.sin(orientation) - c = np.cos(orientation) - return np.array([c, s, -s, c]).reshape(2,2) +def rotation_matrix(orientation): + s = np.sin(orientation) + c = np.cos(orientation) + return np.array([c, s, -s, c]).reshape(2, 2) + + +c = 3e8 # speed of light -c=3e8 # speed of light class Source(object): - def __init__(self,pos): - self.pos=np.array(pos) - assert(self.pos.shape[1]==2) + def __init__(self, pos): + self.pos = np.array(pos) + assert self.pos.shape[1] == 2 - def signal(self,sampling_times): - return np.cos(2*np.pi*sampling_times)+np.sin(2*np.pi*sampling_times)*1j + def signal(self, sampling_times): + return ( + np.cos(2 * np.pi * sampling_times) + np.sin(2 * np.pi * sampling_times) * 1j + ) + + def demod_signal(self, signal, demod_times): + return signal - def demod_signal(self,signal,demod_times): - return signal class IQSource(Source): - def __init__(self,pos,frequency,phase=0,amplitude=1): - super().__init__(pos) - self.frequency=frequency - self.phase=phase - self.amplitude=amplitude + def __init__(self, pos, frequency, phase=0, amplitude=1): + super().__init__(pos) + self.frequency = frequency + self.phase = phase + self.amplitude = amplitude + + def signal(self, sampling_times): + return ( + np.cos(2 * np.pi * sampling_times * self.frequency + self.phase) + + np.sin(2 * np.pi * sampling_times * self.frequency + self.phase) * 1j + ) - def signal(self,sampling_times): - return np.cos(2*np.pi*sampling_times*self.frequency+self.phase)+np.sin(2*np.pi*sampling_times*self.frequency+self.phase)*1j class SinSource(Source): - def __init__(self,pos,frequency,phase=0,amplitude=1): - super().__init__(pos) - self.frequency=frequency - self.phase=phase - self.amplitude=amplitude + def __init__(self, pos, frequency, phase=0, amplitude=1): + super().__init__(pos) + self.frequency = frequency + self.phase = phase + self.amplitude = amplitude + + def signal(self, sampling_times): + # return np.cos(2*np.pi*sampling_times*self.frequency+self.phase)+np.sin(2*np.pi*sampling_times*self.frequency+self.phase)*1j + return self.amplitude * np.sin( + 2 * np.pi * sampling_times * self.frequency + self.phase + ) # .reshape(1,-1) - def signal(self,sampling_times): - #return np.cos(2*np.pi*sampling_times*self.frequency+self.phase)+np.sin(2*np.pi*sampling_times*self.frequency+self.phase)*1j - return (self.amplitude*np.sin(2*np.pi*sampling_times*self.frequency+self.phase))#.reshape(1,-1) class MixedSource(Source): - def __init__(self,source_a,source_b,h=None): - super().__init__(pos) - self.source_a=source_a - self.source_b=source_b - self.h=h + def __init__(self, source_a, source_b, h=None): + super().__init__(pos) + self.source_a = source_a + self.source_b = source_b + self.h = h + + def signal(self, sampling_times): + return self.source_a(sampling_times) * self.source_b(sampling_times) - def signal(self,sampling_times): - return self.source_a(sampling_times)*self.source_b(sampling_times) class NoiseWrapper(Source): - def __init__(self,internal_source,sigma=1): - super().__init__(internal_source.pos) - self.internal_source=internal_source - self.sigma=sigma + def __init__(self, internal_source, sigma=1): + super().__init__(internal_source.pos) + self.internal_source = internal_source + self.sigma = sigma + + def signal(self, sampling_times): + assert sampling_times.ndim == 2 # receivers x time + return ( + self.internal_source.signal(sampling_times) + + ( + np.random.randn(*sampling_times.shape) + + np.random.randn(*sampling_times.shape) * 1j + ) + * self.sigma + ) - def signal(self,sampling_times): - assert(sampling_times.ndim==2) # receivers x time - return self.internal_source.signal(sampling_times) + (np.random.randn(*sampling_times.shape)+np.random.randn(*sampling_times.shape)*1j)*self.sigma class Detector(object): - def __init__(self,sampling_frequency,orientation=0,sigma=0.0): - self.sources=[] - self.source_positions=None - self.receiver_positions=None - self.sampling_frequency=sampling_frequency - self.position_offset=np.zeros(2) - self.orientation=orientation # rotation to the right in radians to apply to receiver array coordinate system - self.sigma=sigma - - def add_source(self,source): - self.sources.append(source) - if self.source_positions is None: - self.source_positions=np.array(source.pos).reshape(1,2) - else: - self.source_positions=np.vstack([ - self.source_positions, - np.array(source.pos).reshape(1,2)]) - - def distance_receiver_to_source(self): - return np.linalg.norm(self.all_receiver_pos()[:,None]-self.source_positions[None],axis=2) - - def rm_sources(self): - self.sources=[] - self.source_positions=None - - def set_receiver_positions(self,receiver_positions): - self.receiver_positions=receiver_positions - - def add_receiver(self,receiver_position): - if self.receiver_positions is None: - self.receiver_positions=np.array(receiver_position).reshape(1,2) - else: - self.receiver_positions=np.vstack([ - self.receiver_positions, - np.array(receiver_position).reshape(1,2)]) - - def n_receivers(self): - return self.receiver_positions.shape[0] - - def all_receiver_pos(self,with_offset=True): - if with_offset: - return self.position_offset+(rotation_matrix(self.orientation) @ self.receiver_positions.T).T - else: - return (rotation_matrix(self.orientation) @ self.receiver_positions.T).T - - def receiver_pos(self,receiver_idx, with_offset=True): - if with_offset: - return self.position_offset+(rotation_matrix(self.orientation) @ self.receiver_positions[receiver_idx].T).T - else: - return (rotation_matrix(self.orientation) @ self.receiver_positions[receiver_idx].T).T - - def get_signal_matrix(self,start_time,duration,rx_lo=0): - n_samples=int(duration*self.sampling_frequency) - base_times=start_time+rf_linspace(0,n_samples-1,n_samples)/self.sampling_frequency - - if self.sigma==0.0: - sample_matrix=np.zeros((self.receiver_positions.shape[0],n_samples),dtype=np.cdouble) # receivers x samples - else: - sample_matrix=np.random.randn( - self.receiver_positions.shape[0],n_samples,2).view(np.cdouble).reshape(self.receiver_positions.shape[0],n_samples)*self.sigma - - if len(self.sources)==0: - return sample_matrix - - distances=self.distance_receiver_to_source().T # sources x receivers # TODO numerical stability, maybe project angle and calculate diff - #TODO diff can be small relative to absolute distance - time_delays=distances/c - base_time_offsets=base_times[None,None]-(distances/c)[...,None] # sources x receivers x sampling intervals - distances_squared=distances**2 - for source_index,_source in enumerate(self.sources): - #get the signal from the source for these times - signal=_source.signal(base_time_offsets[source_index]) #.reshape(base_time_offsets[source_index].shape) # receivers x sampling intervals - normalized_signal=signal/distances_squared[source_index][...,None] - _base_times=np.broadcast_to(base_times,normalized_signal.shape) # broadcast the basetimes for rx_lo on all receivers - demod_times=np.broadcast_to(_base_times.mean(axis=0,keepdims=True),_base_times.shape) #TODO this just takes the average? - ds=_source.demod_signal( - normalized_signal, - demod_times) # TODO nested demod? - sample_matrix+=ds - return sample_matrix #,raw_signal,demod_times,base_time_offsets[0] + def __init__(self, sampling_frequency, orientation=0, sigma=0.0): + self.sources = [] + self.source_positions = None + self.receiver_positions = None + self.sampling_frequency = sampling_frequency + self.position_offset = np.zeros(2) + self.orientation = orientation # rotation to the right in radians to apply to receiver array coordinate system + self.sigma = sigma + + def add_source(self, source): + self.sources.append(source) + if self.source_positions is None: + self.source_positions = np.array(source.pos).reshape(1, 2) + else: + self.source_positions = np.vstack( + [self.source_positions, np.array(source.pos).reshape(1, 2)] + ) + + def distance_receiver_to_source(self): + return np.linalg.norm( + self.all_receiver_pos()[:, None] - self.source_positions[None], axis=2 + ) + + def rm_sources(self): + self.sources = [] + self.source_positions = None + + def set_receiver_positions(self, receiver_positions): + self.receiver_positions = receiver_positions + + def add_receiver(self, receiver_position): + if self.receiver_positions is None: + self.receiver_positions = np.array(receiver_position).reshape(1, 2) + else: + self.receiver_positions = np.vstack( + [self.receiver_positions, np.array(receiver_position).reshape(1, 2)] + ) + + def n_receivers(self): + return self.receiver_positions.shape[0] + + def all_receiver_pos(self, with_offset=True): + if with_offset: + return ( + self.position_offset + + (rotation_matrix(self.orientation) @ self.receiver_positions.T).T + ) + else: + return (rotation_matrix(self.orientation) @ self.receiver_positions.T).T + + def receiver_pos(self, receiver_idx, with_offset=True): + if with_offset: + return ( + self.position_offset + + ( + rotation_matrix(self.orientation) + @ self.receiver_positions[receiver_idx].T + ).T + ) + else: + return ( + rotation_matrix(self.orientation) + @ self.receiver_positions[receiver_idx].T + ).T + + def get_signal_matrix(self, start_time, duration, rx_lo=0): + n_samples = int(duration * self.sampling_frequency) + base_times = ( + start_time + + rf_linspace(0, n_samples - 1, n_samples) / self.sampling_frequency + ) + + if self.sigma == 0.0: + sample_matrix = np.zeros( + (self.receiver_positions.shape[0], n_samples), dtype=np.cdouble + ) # receivers x samples + else: + sample_matrix = ( + np.random.randn(self.receiver_positions.shape[0], n_samples, 2) + .view(np.cdouble) + .reshape(self.receiver_positions.shape[0], n_samples) + * self.sigma + ) + + if len(self.sources) == 0: + return sample_matrix + + distances = ( + self.distance_receiver_to_source().T + ) # sources x receivers # TODO numerical stability, maybe project angle and calculate diff + # TODO diff can be small relative to absolute distance + time_delays = distances / c + base_time_offsets = ( + base_times[None, None] - (distances / c)[..., None] + ) # sources x receivers x sampling intervals + distances_squared = distances**2 + for source_index, _source in enumerate(self.sources): + # get the signal from the source for these times + signal = _source.signal( + base_time_offsets[source_index] + ) # .reshape(base_time_offsets[source_index].shape) # receivers x sampling intervals + normalized_signal = signal / distances_squared[source_index][..., None] + _base_times = np.broadcast_to( + base_times, normalized_signal.shape + ) # broadcast the basetimes for rx_lo on all receivers + demod_times = np.broadcast_to( + _base_times.mean(axis=0, keepdims=True), _base_times.shape + ) # TODO this just takes the average? + ds = _source.demod_signal( + normalized_signal, demod_times + ) # TODO nested demod? + sample_matrix += ds + return sample_matrix # ,raw_signal,demod_times,base_time_offsets[0] @functools.lru_cache(maxsize=1024) -def linear_receiver_positions(n_elements,spacing): - receiver_positions=np.zeros((n_elements,2)) - receiver_positions[:,0]=spacing*(np.arange(n_elements)-(n_elements-1)/2) +def linear_receiver_positions(n_elements, spacing): + receiver_positions = np.zeros((n_elements, 2)) + receiver_positions[:, 0] = spacing * (np.arange(n_elements) - (n_elements - 1) / 2) return receiver_positions + class ULADetector(Detector): - def __init__(self,sampling_frequency,n_elements,spacing,sigma=0.0): - super().__init__(sampling_frequency,sigma=sigma) - self.set_receiver_positions(linear_receiver_positions(n_elements,spacing)) - -@functools.lru_cache(maxsize=1024) -def circular_receiver_positions(n_elements,radius): - theta=(rf_linspace(0,2*np.pi,n_elements+1)[:-1]).reshape(-1,1) - return radius*np.hstack([np.cos(theta),np.sin(theta)]) + def __init__(self, sampling_frequency, n_elements, spacing, sigma=0.0): + super().__init__(sampling_frequency, sigma=sigma) + self.set_receiver_positions(linear_receiver_positions(n_elements, spacing)) -class UCADetector(Detector): - def __init__(self,sampling_frequency,n_elements,radius,sigma=0.0): - super().__init__(sampling_frequency,sigma=sigma) - self.set_receiver_positions(circular_receiver_positions(n_elements,radius)) @functools.lru_cache(maxsize=1024) -def get_thetas(spacing): - thetas=rf_linspace(-np.pi,np.pi,spacing) - return thetas,np.vstack([np.cos(thetas)[None],np.sin(thetas)[None]]).T +def circular_receiver_positions(n_elements, radius): + theta = (rf_linspace(0, 2 * np.pi, n_elements + 1)[:-1]).reshape(-1, 1) + return radius * np.hstack([np.cos(theta), np.sin(theta)]) -if numba: - @jit(nopython=True) - def beamformer_numba_helper(receiver_positions,signal_matrix,carrier_frequency,spacing,thetas,source_vectors): - steer_dot_signal=np.zeros(thetas.shape[0]) - carrier_wavelength=c/carrier_frequency - projection_of_receiver_onto_source_directions=(source_vectors @ receiver_positions.T) - args=2*np.pi*projection_of_receiver_onto_source_directions/carrier_wavelength - steering_vectors=np.exp(-1j*args) - steer_dot_signal=np.absolute(steering_vectors @ signal_matrix).sum(axis=1)/signal_matrix.shape[1] +class UCADetector(Detector): + def __init__(self, sampling_frequency, n_elements, radius, sigma=0.0): + super().__init__(sampling_frequency, sigma=sigma) + self.set_receiver_positions(circular_receiver_positions(n_elements, radius)) + - return thetas,steer_dot_signal,steering_vectors +@functools.lru_cache(maxsize=1024) +def get_thetas(spacing): + thetas = rf_linspace(-np.pi, np.pi, spacing) + return thetas, np.vstack([np.cos(thetas)[None], np.sin(thetas)[None]]).T -def beamformer_numba(receiver_positions,signal_matrix,carrier_frequency,spacing=64+1): - thetas,source_vectors=get_thetas(spacing) - return beamformer_numba_helper(receiver_positions, - signal_matrix, - carrier_frequency, - spacing, - thetas,source_vectors) +if numba: -#from Jon Kraft github + @jit(nopython=True) + def beamformer_numba_helper( + receiver_positions, + signal_matrix, + carrier_frequency, + spacing, + thetas, + source_vectors, + ): + steer_dot_signal = np.zeros(thetas.shape[0]) + carrier_wavelength = c / carrier_frequency + + projection_of_receiver_onto_source_directions = ( + source_vectors @ receiver_positions.T + ) + args = ( + 2 + * np.pi + * projection_of_receiver_onto_source_directions + / carrier_wavelength + ) + steering_vectors = np.exp(-1j * args) + steer_dot_signal = ( + np.absolute(steering_vectors @ signal_matrix).sum(axis=1) + / signal_matrix.shape[1] + ) + + return thetas, steer_dot_signal, steering_vectors + + +def beamformer_numba( + receiver_positions, signal_matrix, carrier_frequency, spacing=64 + 1 +): + thetas, source_vectors = get_thetas(spacing) + return beamformer_numba_helper( + receiver_positions, + signal_matrix, + carrier_frequency, + spacing, + thetas, + source_vectors, + ) + + +# from Jon Kraft github def dbfs(raw_data): # function to convert IQ samples to FFT plot, scaled in dBFS NumSamples = len(raw_data) @@ -220,12 +310,14 @@ def dbfs(raw_data): y = raw_data * win s_fft = np.fft.fft(y) / np.sum(win) s_shift = np.fft.fftshift(s_fft) - s_dbfs = 20*np.log10(np.abs(s_shift)/(2**11)) # Pluto is a signed 12 bit ADC, so use 2^11 to convert to dBFS + s_dbfs = 20 * np.log10( + np.abs(s_shift) / (2**11) + ) # Pluto is a signed 12 bit ADC, so use 2^11 to convert to dBFS return s_dbfs ### -''' +""" Beamformer takes as input the receiver positions signal marix representing signal received at those positions @@ -247,21 +339,37 @@ def dbfs(raw_data): 0 -> x=0, y=1 pi/2 -> x=1, y=0 -pi/2 -> x=-1, y=0 -''' -### -def beamformer(receiver_positions,signal_matrix,carrier_frequency,calibration=None,spacing=64+1,offset=0.0): - thetas=np.linspace(-np.pi,np.pi,spacing)#-offset - source_vectors=np.vstack([np.sin(thetas+offset)[None],np.cos(thetas+offset)[None]]).T +""" - projection_of_receiver_onto_source_directions=(source_vectors @ receiver_positions.T) - carrier_wavelength=c/carrier_frequency - args=2*np.pi*projection_of_receiver_onto_source_directions/carrier_wavelength - steering_vectors=np.exp(-1j*args) +### +def beamformer( + receiver_positions, + signal_matrix, + carrier_frequency, + calibration=None, + spacing=64 + 1, + offset=0.0, +): + thetas = np.linspace(-np.pi, np.pi, spacing) # -offset + source_vectors = np.vstack( + [np.sin(thetas + offset)[None], np.cos(thetas + offset)[None]] + ).T + + projection_of_receiver_onto_source_directions = ( + source_vectors @ receiver_positions.T + ) + + carrier_wavelength = c / carrier_frequency + args = ( + 2 * np.pi * projection_of_receiver_onto_source_directions / carrier_wavelength + ) + steering_vectors = np.exp(-1j * args) if calibration is not None: - steering_vectors=steering_vectors*calibration[None] - #the delay sum is performed in the matmul step, the absolute is over the summed value - phase_adjusted=np.matmul(steering_vectors,signal_matrix) # this is adjust and sum in one step - steer_dot_signal=np.absolute(phase_adjusted).mean(axis=1) # mean over samples - return thetas,steer_dot_signal,steering_vectors - + steering_vectors = steering_vectors * calibration[None] + # the delay sum is performed in the matmul step, the absolute is over the summed value + phase_adjusted = np.matmul( + steering_vectors, signal_matrix + ) # this is adjust and sum in one step + steer_dot_signal = np.absolute(phase_adjusted).mean(axis=1) # mean over samples + return thetas, steer_dot_signal, steering_vectors diff --git a/software/model_training_and_inference/utils/save_load.py b/software/model_training_and_inference/utils/save_load.py index 75bd39d4..1e940bf8 100644 --- a/software/model_training_and_inference/utils/save_load.py +++ b/software/model_training_and_inference/utils/save_load.py @@ -2,10 +2,11 @@ import io import torch -#borrowed from online + +# borrowed from online class CPU_Unpickler(pickle.Unpickler): - def find_class(self, module, name): - if module == 'torch.storage' and name == '_load_from_bytes': - return lambda b: torch.load(io.BytesIO(b), map_location='cpu') - else: - return super().find_class(module, name) + def find_class(self, module, name): + if module == "torch.storage" and name == "_load_from_bytes": + return lambda b: torch.load(io.BytesIO(b), map_location="cpu") + else: + return super().find_class(module, name) diff --git a/software/model_training_and_inference/utils/spf_dataset.py b/software/model_training_and_inference/utils/spf_dataset.py index 337b6710..e1204ff5 100644 --- a/software/model_training_and_inference/utils/spf_dataset.py +++ b/software/model_training_and_inference/utils/spf_dataset.py @@ -11,408 +11,575 @@ import torch from torch.utils.data import DataLoader, Dataset -from utils.image_utils import (detector_positions_to_theta_grid, - labels_to_source_images, radio_to_image) +from utils.image_utils import ( + detector_positions_to_theta_grid, + labels_to_source_images, + radio_to_image, +) from utils.spf_generate import generate_session from compress_pickle import dump, load -output_cols={ # maybe this should get moved to the dataset part... - 'src_pos':[0,1], - 'src_theta':[2], - 'src_dist':[3], - 'det_delta':[4,5], - 'det_theta':[6], - 'det_space':[7], - 'src_v':[8,9] +output_cols = { # maybe this should get moved to the dataset part... + "src_pos": [0, 1], + "src_theta": [2], + "src_dist": [3], + "det_delta": [4, 5], + "det_theta": [6], + "det_space": [7], + "src_v": [8, 9], } -input_cols={ - 'det_pos':[0,1], - 'time':[2], - 'space_delta':[3,4], - 'space_theta':[5], - 'space_dist':[6], - 'det_theta2':[7], +input_cols = { + "det_pos": [0, 1], + "time": [2], + "space_delta": [3, 4], + "space_theta": [5], + "space_dist": [6], + "det_theta2": [7], } -#from stackoverflow + +# from stackoverflow class dotdict(dict): """dot.notation access to dictionary attributes""" + __getattr__ = dict.get __setattr__ = dict.__setitem__ __delattr__ = dict.__delitem__ + class SessionsDataset(Dataset): + def __init__(self, root_dir, snapshots_in_sample=5): + """ + Arguments: + root_dir (string): Directory with all the images. + """ + self.root_dir = root_dir + self.args = load("/".join([self.root_dir, "args.pkl"]), compression="lzma") + assert self.args.time_steps >= snapshots_in_sample + self.samples_per_session = self.args.time_steps - snapshots_in_sample + 1 + self.snapshots_in_sample = snapshots_in_sample + if not self.args.live: + print("NOT LIVE") + self.filenames = sorted( + filter( + lambda x: "args" not in x, + ["%s/%s" % (self.root_dir, x) for x in os.listdir(self.root_dir)], + ) + ) + if self.args.sessions != len( + self.filenames + ): # make sure its the right dataset + print("WARNING DATASET LOOKS LIKE IT IS MISSING SOME SESSIONS!") + + def idx_to_filename_and_start_idx(self, idx): + assert idx >= 0 and idx <= self.samples_per_session * len(self.filenames) + return ( + self.filenames[idx // self.samples_per_session], + idx % self.samples_per_session, + ) + + def __len__(self): + if self.args.live: + return self.samples_per_session * self.args.sessions + else: + return self.samples_per_session * len(self.filenames) + + def __getitem__(self, idx): + session_idx = idx // self.samples_per_session + start_idx = idx % self.samples_per_session + + if self.args.live: + session = generate_session((self.args, session_idx)) + else: + session = load(self.filenames[session_idx], compression="lzma") + end_idx = start_idx + self.snapshots_in_sample + return {k: session[k][start_idx:end_idx] for k in session.keys()} - def __init__(self, root_dir, snapshots_in_sample=5): - """ - Arguments: - root_dir (string): Directory with all the images. - """ - self.root_dir = root_dir - self.args=load("/".join([self.root_dir,'args.pkl']),compression="lzma") - assert(self.args.time_steps>=snapshots_in_sample) - self.samples_per_session=self.args.time_steps-snapshots_in_sample+1 - self.snapshots_in_sample=snapshots_in_sample - if not self.args.live: - print("NOT LIVE") - self.filenames=sorted(filter(lambda x : 'args' not in x ,[ "%s/%s" % (self.root_dir,x) for x in os.listdir(self.root_dir)])) - if self.args.sessions!=len(self.filenames): # make sure its the right dataset - print("WARNING DATASET LOOKS LIKE IT IS MISSING SOME SESSIONS!") - - def idx_to_filename_and_start_idx(self,idx): - assert(idx>=0 and idx<=self.samples_per_session*len(self.filenames)) - return self.filenames[idx//self.samples_per_session],idx%self.samples_per_session - - def __len__(self): - if self.args.live: - return self.samples_per_session*self.args.sessions - else: - return self.samples_per_session*len(self.filenames) - - def __getitem__(self, idx): - session_idx=idx//self.samples_per_session - start_idx=idx%self.samples_per_session - - if self.args.live: - session=generate_session((self.args,session_idx)) - else: - session=load(self.filenames[session_idx],compression="lzma") - end_idx=start_idx+self.snapshots_in_sample - return { k:session[k][start_idx:end_idx] for k in session.keys()} class SessionsDatasetReal(Dataset): - - def get_m(self,filename,bypass=False): - if bypass: - return np.memmap( - filename, - dtype='float32', - mode='r', - shape=(self.snapshots_in_file,self.nthetas+5)) - - if filename not in self.m_cache: - self.m_cache[filename]=np.memmap( - filename, - dtype='float32', - mode='r', - shape=(self.snapshots_in_file,self.nthetas+5)) - return self.m_cache[filename] - - def check_file(self,filename): - try: - m=self.get_m(filename,bypass=True) - except: - print("DROP",filename) - return False - status=not (np.abs(m).mean(axis=1)==0).any() - if status==False: - print("DROP",filename) - return status - - def __init__(self, - root_dir, - snapshots_in_file=400000, - nthetas=65, - snapshots_in_sample=128, - nsources=1, - width=3000, - receiver_pos_x=1352.7, - receiver_pos_y=2647.7, - receiver_spacing=60.0, - step_size=1, - seed=1337 - ): - #time_step,x,y,mean_angle,_mean_angle #0,1,2,3,4 - #m = np.memmap(filename, dtype='float32', mode='r', shape=(,70)) - """ - Arguments: - root_dir (string): Directory with all the images. - """ - assert(nsources==1) #TODO implement more - self.root_dir = root_dir - self.nthetas=nthetas - self.thetas=np.linspace(-np.pi,np.pi,self.nthetas) - self.args=dotdict({ - 'width':width, - }) - self.m_cache={} - self.receiver_pos=np.array([ - [receiver_pos_x-receiver_spacing/2,receiver_pos_y], - [receiver_pos_x+receiver_spacing/2,receiver_pos_y] - ]) - self.detector_position=np.array([[receiver_pos_x,receiver_pos_y]]) - self.snapshots_in_file=snapshots_in_file - self.snapshots_in_sample=snapshots_in_sample - self.step_size=step_size - self.filenames=sorted( - filter( self.check_file, - filter( - lambda x : '.npy' in x ,[ "%s/%s" % (self.root_dir,x) for x in os.listdir(self.root_dir)]))) - self.rng = np.random.default_rng(seed) - self.rng.shuffle(self.filenames) - #self.datas=[ - # np.memmap( - # filename, - # dtype='float32', - # mode='r', - # shape=(self.snapshots_in_file,self.nthetas+5)) for filename in self.filenames - #] - self.samples_per_file=[ - self.get_m(filename,bypass=True).shape[0]-(self.snapshots_in_sample*self.step_size) for filename in self.filenames - ] - self.cumsum_samples_per_file=np.cumsum([0]+self.samples_per_file) - self.len=sum(self.samples_per_file) - self.zeros=np.zeros((self.snapshots_in_sample,5)) - self.ones=np.ones((self.snapshots_in_sample,5)) - self.widths=np.ones((self.snapshots_in_sample,1),dtype=np.int32)*self.args.width - self.halfpis=-np.ones((self.snapshots_in_sample,1))*np.pi/2 - idx_to_fileidx_and_sampleidx={} - #print("WARNING BY DEFAULT FLIPPING RADIO FEATURE SINCE COORDS WERE WRONG IN PI COLLECT!") - - def idx_to_fileidx_and_startidx(self,idx): - file_idx=bisect.bisect_right(self.cumsum_samples_per_file, idx)-1 - if file_idx>=len(self.samples_per_file): - return None - start_idx=(idx-self.cumsum_samples_per_file[file_idx]) #*(self.snapshots_in_sample*self.step_size) - return file_idx,start_idx - - def __len__(self): - return self.len - - def __getitem__(self, idx): - fileidx,startidx=self.idx_to_fileidx_and_startidx(idx) - m=self.get_m(self.filenames[fileidx])[startidx:startidx+self.snapshots_in_sample*self.step_size:self.step_size] - - detector_position_at_t=np.broadcast_to(self.detector_position, (m.shape[0], 2)) - detector_orientation_at_t=self.zeros[:,[0]] - source_positions_at_t=m[:,1:3][:,None] - diffs=source_positions_at_t-detector_position_at_t[:,None] # broadcast over nsource dimension - #diffs=(batchsize, nsources, 2) - source_theta_at_t=np.arctan2(diffs[:,0,[0]],diffs[:,0,[1]]) # rotation to the right around x=0, y+ - return { - 'broadcasting_positions_at_t':self.ones[:,[0]][:,None], # TODO multi source - 'source_positions_at_t':source_positions_at_t, - 'source_velocities_at_t':self.zeros[:,:2][:,None], #TODO calc velocity, - 'receiver_positions_at_t':np.broadcast_to(self.receiver_pos[None], (m.shape[0],2, 2)), - 'beam_former_outputs_at_t':m[:,5:], - 'thetas_at_t':np.broadcast_to(self.thetas[None], (m.shape[0],self.thetas.shape[0])), - 'time_stamps':m[:,[0]], - 'width_at_t': self.widths, - 'detector_orientation_at_t':detector_orientation_at_t, #self.halfpis*0,#np.arctan2(1,0)=np.pi/2 - 'detector_position_at_t':detector_position_at_t, - 'source_theta_at_t':source_theta_at_t, - 'source_distance_at_t':self.zeros[:,[0]][:,None], - } + def get_m(self, filename, bypass=False): + if bypass: + return np.memmap( + filename, + dtype="float32", + mode="r", + shape=(self.snapshots_in_file, self.nthetas + 5), + ) + + if filename not in self.m_cache: + self.m_cache[filename] = np.memmap( + filename, + dtype="float32", + mode="r", + shape=(self.snapshots_in_file, self.nthetas + 5), + ) + return self.m_cache[filename] + + def check_file(self, filename): + try: + m = self.get_m(filename, bypass=True) + except: + print("DROP", filename) + return False + status = not (np.abs(m).mean(axis=1) == 0).any() + if status == False: + print("DROP", filename) + return status + + def __init__( + self, + root_dir, + snapshots_in_file=400000, + nthetas=65, + snapshots_in_sample=128, + nsources=1, + width=3000, + receiver_pos_x=1352.7, + receiver_pos_y=2647.7, + receiver_spacing=60.0, + step_size=1, + seed=1337, + ): + # time_step,x,y,mean_angle,_mean_angle #0,1,2,3,4 + # m = np.memmap(filename, dtype='float32', mode='r', shape=(,70)) + """ + Arguments: + root_dir (string): Directory with all the images. + """ + assert nsources == 1 # TODO implement more + self.root_dir = root_dir + self.nthetas = nthetas + self.thetas = np.linspace(-np.pi, np.pi, self.nthetas) + self.args = dotdict( + { + "width": width, + } + ) + self.m_cache = {} + self.receiver_pos = np.array( + [ + [receiver_pos_x - receiver_spacing / 2, receiver_pos_y], + [receiver_pos_x + receiver_spacing / 2, receiver_pos_y], + ] + ) + self.detector_position = np.array([[receiver_pos_x, receiver_pos_y]]) + self.snapshots_in_file = snapshots_in_file + self.snapshots_in_sample = snapshots_in_sample + self.step_size = step_size + self.filenames = sorted( + filter( + self.check_file, + filter( + lambda x: ".npy" in x, + ["%s/%s" % (self.root_dir, x) for x in os.listdir(self.root_dir)], + ), + ) + ) + self.rng = np.random.default_rng(seed) + self.rng.shuffle(self.filenames) + # self.datas=[ + # np.memmap( + # filename, + # dtype='float32', + # mode='r', + # shape=(self.snapshots_in_file,self.nthetas+5)) for filename in self.filenames + # ] + self.samples_per_file = [ + self.get_m(filename, bypass=True).shape[0] + - (self.snapshots_in_sample * self.step_size) + for filename in self.filenames + ] + self.cumsum_samples_per_file = np.cumsum([0] + self.samples_per_file) + self.len = sum(self.samples_per_file) + self.zeros = np.zeros((self.snapshots_in_sample, 5)) + self.ones = np.ones((self.snapshots_in_sample, 5)) + self.widths = ( + np.ones((self.snapshots_in_sample, 1), dtype=np.int32) * self.args.width + ) + self.halfpis = -np.ones((self.snapshots_in_sample, 1)) * np.pi / 2 + idx_to_fileidx_and_sampleidx = {} + # print("WARNING BY DEFAULT FLIPPING RADIO FEATURE SINCE COORDS WERE WRONG IN PI COLLECT!") + + def idx_to_fileidx_and_startidx(self, idx): + file_idx = bisect.bisect_right(self.cumsum_samples_per_file, idx) - 1 + if file_idx >= len(self.samples_per_file): + return None + start_idx = ( + idx - self.cumsum_samples_per_file[file_idx] + ) # *(self.snapshots_in_sample*self.step_size) + return file_idx, start_idx + + def __len__(self): + return self.len + + def __getitem__(self, idx): + fileidx, startidx = self.idx_to_fileidx_and_startidx(idx) + m = self.get_m(self.filenames[fileidx])[ + startidx : startidx + + self.snapshots_in_sample * self.step_size : self.step_size + ] + + detector_position_at_t = np.broadcast_to( + self.detector_position, (m.shape[0], 2) + ) + detector_orientation_at_t = self.zeros[:, [0]] + source_positions_at_t = m[:, 1:3][:, None] + diffs = ( + source_positions_at_t - detector_position_at_t[:, None] + ) # broadcast over nsource dimension + # diffs=(batchsize, nsources, 2) + source_theta_at_t = np.arctan2( + diffs[:, 0, [0]], diffs[:, 0, [1]] + ) # rotation to the right around x=0, y+ + return { + "broadcasting_positions_at_t": self.ones[:, [0]][ + :, None + ], # TODO multi source + "source_positions_at_t": source_positions_at_t, + "source_velocities_at_t": self.zeros[:, :2][:, None], # TODO calc velocity, + "receiver_positions_at_t": np.broadcast_to( + self.receiver_pos[None], (m.shape[0], 2, 2) + ), + "beam_former_outputs_at_t": m[:, 5:], + "thetas_at_t": np.broadcast_to( + self.thetas[None], (m.shape[0], self.thetas.shape[0]) + ), + "time_stamps": m[:, [0]], + "width_at_t": self.widths, + "detector_orientation_at_t": detector_orientation_at_t, # self.halfpis*0,#np.arctan2(1,0)=np.pi/2 + "detector_position_at_t": detector_position_at_t, + "source_theta_at_t": source_theta_at_t, + "source_distance_at_t": self.zeros[:, [0]][:, None], + } class SessionsDatasetTask1(SessionsDataset): - def __getitem__(self,idx): - d=super().__getitem__(idx) - #featurie a really simple way - x=torch.Tensor(np.hstack([ - d['receiver_positions_at_t'].reshape(self.snapshots_in_sample,-1), - d['beam_former_outputs_at_t'].reshape(self.snapshots_in_sample,-1), - #d['signal_matrixs'].reshape(self.snapshots_in_sample,-1) - d['time_stamps'].reshape(self.snapshots_in_sample,-1)-d['time_stamps'][0], - ])) - y=torch.Tensor(d['source_positions_at_t'][:,0]) - return x,y + def __getitem__(self, idx): + d = super().__getitem__(idx) + # featurie a really simple way + x = torch.Tensor( + np.hstack( + [ + d["receiver_positions_at_t"].reshape(self.snapshots_in_sample, -1), + d["beam_former_outputs_at_t"].reshape(self.snapshots_in_sample, -1), + # d['signal_matrixs'].reshape(self.snapshots_in_sample,-1) + d["time_stamps"].reshape(self.snapshots_in_sample, -1) + - d["time_stamps"][0], + ] + ) + ) + y = torch.Tensor(d["source_positions_at_t"][:, 0]) + return x, y + +def pos_to_rel(p, width): + return 2 * (p / width - 0.5) -def pos_to_rel(p,width): - return 2*(p/width-0.5) -def rel_to_pos(r,width): - return ((r/2)+0.5)*width +def rel_to_pos(r, width): + return ((r / 2) + 0.5) * width + class SessionsDatasetRealTask2(SessionsDatasetReal): - def __getitem__(self,idx): - d=super().__getitem__(idx) - #normalize these before heading out - d['source_positions_at_t_normalized_centered']=2*(d['source_positions_at_t']/self.args.width-0.5) - d['source_velocities_at_t_normalized']=d['source_velocities_at_t']/self.args.width - d['detector_position_at_t_normalized_centered']=2*(d['detector_position_at_t']/self.args.width-0.5) - d['source_distance_at_t_normalized']=d['source_distance_at_t'].mean(axis=2)/(self.args.width/2) - return d #,d['source_positions_at_t'] + def __getitem__(self, idx): + d = super().__getitem__(idx) + # normalize these before heading out + d["source_positions_at_t_normalized_centered"] = 2 * ( + d["source_positions_at_t"] / self.args.width - 0.5 + ) + d["source_velocities_at_t_normalized"] = ( + d["source_velocities_at_t"] / self.args.width + ) + d["detector_position_at_t_normalized_centered"] = 2 * ( + d["detector_position_at_t"] / self.args.width - 0.5 + ) + d["source_distance_at_t_normalized"] = d["source_distance_at_t"].mean( + axis=2 + ) / (self.args.width / 2) + return d # ,d['source_positions_at_t'] -class SessionsDatasetTask2(SessionsDataset): - def __getitem__(self,idx): - d=super().__getitem__(idx) - #normalize these before heading out - d['source_positions_at_t_normalized_centered']=2*(d['source_positions_at_t']/self.args.width-0.5) - d['source_velocities_at_t_normalized']=d['source_velocities_at_t']/self.args.width - d['detector_position_at_t_normalized_centered']=2*(d['detector_position_at_t']/self.args.width-0.5) - d['source_distance_at_t_normalized']=d['source_distance_at_t'].mean(axis=2)/(self.args.width/2) - return d #,d['source_positions_at_t'] -class SessionsDatasetTask2WithImages(SessionsDataset): - def __getitem__(self,idx): - d=super().__getitem__(idx) - #normalize these before heading out - d['source_positions_at_t_normalized']=2*(d['source_positions_at_t']/self.args.width-0.5) - d['source_velocities_at_t_normalized']=d['source_velocities_at_t']/self.args.width - d['detector_position_at_t_normalized']=2*(d['detector_position_at_t']/self.args.width-0.5) - d['source_distance_at_t_normalized']=d['source_distance_at_t'].mean(axis=2)/(self.args.width/2) +class SessionsDatasetTask2(SessionsDataset): + def __getitem__(self, idx): + d = super().__getitem__(idx) + # normalize these before heading out + d["source_positions_at_t_normalized_centered"] = 2 * ( + d["source_positions_at_t"] / self.args.width - 0.5 + ) + d["source_velocities_at_t_normalized"] = ( + d["source_velocities_at_t"] / self.args.width + ) + d["detector_position_at_t_normalized_centered"] = 2 * ( + d["detector_position_at_t"] / self.args.width - 0.5 + ) + d["source_distance_at_t_normalized"] = d["source_distance_at_t"].mean( + axis=2 + ) / (self.args.width / 2) + return d # ,d['source_positions_at_t'] - d['source_image_at_t']=labels_to_source_images(d['source_positions_at_t'][None],self.args.width)[0] - d['detector_theta_image_at_t']=detector_positions_to_theta_grid(d['detector_position_at_t'][None],self.args.width)[0] - d['radio_image_at_t']=radio_to_image(d['beam_former_outputs_at_t'][None],d['detector_theta_image_at_t'][None],d['detector_orientation_at_t'][None])[0] - return d #,d['source_positions_at_t'] +class SessionsDatasetTask2WithImages(SessionsDataset): + def __getitem__(self, idx): + d = super().__getitem__(idx) + # normalize these before heading out + d["source_positions_at_t_normalized"] = 2 * ( + d["source_positions_at_t"] / self.args.width - 0.5 + ) + d["source_velocities_at_t_normalized"] = ( + d["source_velocities_at_t"] / self.args.width + ) + d["detector_position_at_t_normalized"] = 2 * ( + d["detector_position_at_t"] / self.args.width - 0.5 + ) + d["source_distance_at_t_normalized"] = d["source_distance_at_t"].mean( + axis=2 + ) / (self.args.width / 2) + + d["source_image_at_t"] = labels_to_source_images( + d["source_positions_at_t"][None], self.args.width + )[0] + d["detector_theta_image_at_t"] = detector_positions_to_theta_grid( + d["detector_position_at_t"][None], self.args.width + )[0] + d["radio_image_at_t"] = radio_to_image( + d["beam_former_outputs_at_t"][None], + d["detector_theta_image_at_t"][None], + d["detector_orientation_at_t"][None], + )[0] + + return d # ,d['source_positions_at_t'] def collate_fn_beamformer(_in): - d={ k:torch.from_numpy(np.stack([ x[k] for x in _in ])) for k in _in[0]} - b,s,n_sources,_=d['source_positions_at_t'].shape - - times=d['time_stamps']/(0.00001+d['time_stamps'].max(axis=2,keepdim=True)[0]) - - source_theta=d['source_theta_at_t'].mean(axis=2) - distances=d['source_distance_at_t_normalized'].mean(axis=2,keepdims=True) - _,_,beam_former_bins=d['beam_former_outputs_at_t'].shape - perfect_labels=torch.zeros((b,s,beam_former_bins)) - - idxs=(beam_former_bins*(d['source_theta_at_t']+np.pi)/(2*np.pi)).int() - smooth_bins=int(beam_former_bins*0.25*0.5) - for _b in torch.arange(b): - for _s in torch.arange(s): - for smooth in range(-smooth_bins,smooth_bins+1): - perfect_labels[_b,_s,(idxs[_b,_s]+smooth)%beam_former_bins]=1/(1+smooth**2) - perfect_labels[_b,_s]/=perfect_labels[_b,_s].sum()+1e-9 - r= {'input':torch.concatenate([ - #d['signal_matrixs_at_t'].reshape(b,s,-1), - (d['signal_matrixs_at_t']/d['signal_matrixs_at_t'].abs().mean(axis=[2,3],keepdims=True)).reshape(b,s,-1), # normalize the data - d['signal_matrixs_at_t'].abs().mean(axis=[2,3],keepdims=True).reshape(b,s,-1), # - d['detector_orientation_at_t'].to(torch.complex64)], - axis=2), - 'beamformer':d['beam_former_outputs_at_t'], - 'labels':perfect_labels, - 'thetas':source_theta} - return r + d = {k: torch.from_numpy(np.stack([x[k] for x in _in])) for k in _in[0]} + b, s, n_sources, _ = d["source_positions_at_t"].shape + + times = d["time_stamps"] / (0.00001 + d["time_stamps"].max(axis=2, keepdim=True)[0]) + + source_theta = d["source_theta_at_t"].mean(axis=2) + distances = d["source_distance_at_t_normalized"].mean(axis=2, keepdims=True) + _, _, beam_former_bins = d["beam_former_outputs_at_t"].shape + perfect_labels = torch.zeros((b, s, beam_former_bins)) + + idxs = (beam_former_bins * (d["source_theta_at_t"] + np.pi) / (2 * np.pi)).int() + smooth_bins = int(beam_former_bins * 0.25 * 0.5) + for _b in torch.arange(b): + for _s in torch.arange(s): + for smooth in range(-smooth_bins, smooth_bins + 1): + perfect_labels[ + _b, _s, (idxs[_b, _s] + smooth) % beam_former_bins + ] = 1 / (1 + smooth**2) + perfect_labels[_b, _s] /= perfect_labels[_b, _s].sum() + 1e-9 + r = { + "input": torch.concatenate( + [ + # d['signal_matrixs_at_t'].reshape(b,s,-1), + ( + d["signal_matrixs_at_t"] + / d["signal_matrixs_at_t"].abs().mean(axis=[2, 3], keepdims=True) + ).reshape( + b, s, -1 + ), # normalize the data + d["signal_matrixs_at_t"] + .abs() + .mean(axis=[2, 3], keepdims=True) + .reshape(b, s, -1), # + d["detector_orientation_at_t"].to(torch.complex64), + ], + axis=2, + ), + "beamformer": d["beam_former_outputs_at_t"], + "labels": perfect_labels, + "thetas": source_theta, + } + return r + def collate_fn_transformer_filter(_in): - d={ k:torch.from_numpy(np.stack([ x[k] for x in _in ])) for k in _in[0]} - b,s,n_sources,_=d['source_positions_at_t'].shape - - normalized_01_times=d['time_stamps']/(0.0000001+d['time_stamps'].max(axis=1,keepdim=True)[0]) - normalized_times=(d['time_stamps']-d['time_stamps'].max(axis=1,keepdims=True)[0])/100 - - normalized_pirads_detector_theta=d['detector_orientation_at_t']/np.pi - - space_diffs=(d['detector_position_at_t_normalized_centered'][:,:-1]-d['detector_position_at_t_normalized_centered'][:,1:]) - space_delta=torch.cat([ - torch.zeros(b,1,2), - space_diffs, - ],axis=1) - - normalized_pirads_space_theta=torch.cat([ - torch.zeros(b,1,1), - (torch.atan2(space_diffs[...,1],space_diffs[...,0]))[:,:,None]/np.pi - ],axis=1) - - space_dist=torch.cat([ - torch.zeros(b,1,1), - torch.sqrt(torch.pow(space_diffs,2).sum(axis=2,keepdim=True)) - ],axis=1) - return { - 'drone_state':torch.cat( - [ - d['detector_position_at_t_normalized_centered'], # 2: 2 - #normalized_01_times, #-times.max(axis=2,keepdim=True)[0], # 1: 3 - #normalized_times, #-times.max(axis=2,keepdim=True)[0], # 1: 3 - space_delta, # 2: 4 - normalized_pirads_space_theta, # 1: 5 - space_dist, #1: 6 - normalized_pirads_detector_theta, #1: 7 - ],dim=2).float(), - 'times':normalized_times, - 'emitter_position_and_velocity':torch.cat([ - d['source_positions_at_t_normalized_centered'], - d['source_velocities_at_t_normalized'], - ],dim=3), - 'emitters_broadcasting':d['broadcasting_positions_at_t'], - 'emitters_n_broadcasts':d['broadcasting_positions_at_t'].cumsum(axis=1), - 'radio_feature':torch.cat( - [ - torch.log(d['beam_former_outputs_at_t'].mean(axis=2,keepdim=True))/20, - d['beam_former_outputs_at_t']/d['beam_former_outputs_at_t'].mean(axis=2,keepdim=True), # maybe pass in log values? - ], - dim=2 - ).float() - - } + d = {k: torch.from_numpy(np.stack([x[k] for x in _in])) for k in _in[0]} + b, s, n_sources, _ = d["source_positions_at_t"].shape + + normalized_01_times = d["time_stamps"] / ( + 0.0000001 + d["time_stamps"].max(axis=1, keepdim=True)[0] + ) + normalized_times = ( + d["time_stamps"] - d["time_stamps"].max(axis=1, keepdims=True)[0] + ) / 100 + + normalized_pirads_detector_theta = d["detector_orientation_at_t"] / np.pi + + space_diffs = ( + d["detector_position_at_t_normalized_centered"][:, :-1] + - d["detector_position_at_t_normalized_centered"][:, 1:] + ) + space_delta = torch.cat( + [ + torch.zeros(b, 1, 2), + space_diffs, + ], + axis=1, + ) + + normalized_pirads_space_theta = torch.cat( + [ + torch.zeros(b, 1, 1), + (torch.atan2(space_diffs[..., 1], space_diffs[..., 0]))[:, :, None] / np.pi, + ], + axis=1, + ) + + space_dist = torch.cat( + [ + torch.zeros(b, 1, 1), + torch.sqrt(torch.pow(space_diffs, 2).sum(axis=2, keepdim=True)), + ], + axis=1, + ) + return { + "drone_state": torch.cat( + [ + d["detector_position_at_t_normalized_centered"], # 2: 2 + # normalized_01_times, #-times.max(axis=2,keepdim=True)[0], # 1: 3 + # normalized_times, #-times.max(axis=2,keepdim=True)[0], # 1: 3 + space_delta, # 2: 4 + normalized_pirads_space_theta, # 1: 5 + space_dist, # 1: 6 + normalized_pirads_detector_theta, # 1: 7 + ], + dim=2, + ).float(), + "times": normalized_times, + "emitter_position_and_velocity": torch.cat( + [ + d["source_positions_at_t_normalized_centered"], + d["source_velocities_at_t_normalized"], + ], + dim=3, + ), + "emitters_broadcasting": d["broadcasting_positions_at_t"], + "emitters_n_broadcasts": d["broadcasting_positions_at_t"].cumsum(axis=1), + "radio_feature": torch.cat( + [ + torch.log(d["beam_former_outputs_at_t"].mean(axis=2, keepdim=True)) + / 20, + d["beam_former_outputs_at_t"] + / d["beam_former_outputs_at_t"].mean( + axis=2, keepdim=True + ), # maybe pass in log values? + ], + dim=2, + ).float(), + } + def collate_fn(_in): - d={ k:torch.from_numpy(np.stack([ x[k] for x in _in ])) for k in _in[0]} - b,s,n_sources,_=d['source_positions_at_t'].shape - - #times=d['time_stamps']/(0.00001+d['time_stamps'].max(axis=2,keepdim=True)[0]) - times=d['time_stamps']/(0.0000001+d['time_stamps'].max(axis=1,keepdim=True)[0]) - - #deal with source positions - source_position=d['source_positions_at_t_normalized'][torch.where(d['broadcasting_positions_at_t']==1)[:-1]].reshape(b,s,2).float() - source_velocity=d['source_velocities_at_t_normalized'][torch.where(d['broadcasting_positions_at_t']==1)[:-1]].reshape(b,s,2).float() - - #deal with detector position features - #diffs=source_positions-d['detector_position_at_t_normalized'] - #source_thetab=(torch.atan2(diffs[...,1],diffs[...,0]))[:,:,None]/np.pi # batch, snapshot,1, x ,y - source_theta=d['source_theta_at_t'].mean(axis=2)/np.pi - detector_theta=d['detector_orientation_at_t']/np.pi - #distancesb=torch.sqrt(torch.pow(diffs, 2).sum(axis=2,keepdim=True)) - distances=d['source_distance_at_t_normalized'].mean(axis=2,keepdims=True) - space_diffs=(d['detector_position_at_t_normalized'][:,:-1]-d['detector_position_at_t_normalized'][:,1:]) - space_delta=torch.cat([ - torch.zeros(b,1,2), - space_diffs, - ],axis=1) - - space_theta=torch.cat([ - torch.zeros(b,1,1), - (torch.atan2(space_diffs[...,1],space_diffs[...,0]))[:,:,None]/np.pi - ],axis=1) - - space_dist=torch.cat([ - torch.zeros(b,1,1), - torch.sqrt(torch.pow(space_diffs,2).sum(axis=2,keepdim=True)) - ],axis=1) - - #create the labels - labels=torch.cat([ - source_position, # zero center the positions - (source_theta+detector_theta+1)%2.0-1, # initialy in units of np.pi? - distances, # try to zero center? - space_delta, - space_theta, - space_dist, - source_velocity, - ], axis=2).float() #.to(device) - #breakpoint() - #create the features - inputs={ - 'drone_state':torch.cat( - [ - d['detector_position_at_t_normalized'], - times, #-times.max(axis=2,keepdim=True)[0], - space_delta, - space_theta, - space_dist, - detector_theta, - ],dim=2).float(), - 'radio_feature':torch.cat( - [ - torch.log(d['beam_former_outputs_at_t'].mean(axis=2,keepdim=True))/20, - d['beam_former_outputs_at_t']/d['beam_former_outputs_at_t'].mean(axis=2,keepdim=True), # maybe pass in log values? - ], - dim=2 - ).float() + d = {k: torch.from_numpy(np.stack([x[k] for x in _in])) for k in _in[0]} + b, s, n_sources, _ = d["source_positions_at_t"].shape + + # times=d['time_stamps']/(0.00001+d['time_stamps'].max(axis=2,keepdim=True)[0]) + times = d["time_stamps"] / ( + 0.0000001 + d["time_stamps"].max(axis=1, keepdim=True)[0] + ) + + # deal with source positions + source_position = ( + d["source_positions_at_t_normalized"][ + torch.where(d["broadcasting_positions_at_t"] == 1)[:-1] + ] + .reshape(b, s, 2) + .float() + ) + source_velocity = ( + d["source_velocities_at_t_normalized"][ + torch.where(d["broadcasting_positions_at_t"] == 1)[:-1] + ] + .reshape(b, s, 2) + .float() + ) + + # deal with detector position features + # diffs=source_positions-d['detector_position_at_t_normalized'] + # source_thetab=(torch.atan2(diffs[...,1],diffs[...,0]))[:,:,None]/np.pi # batch, snapshot,1, x ,y + source_theta = d["source_theta_at_t"].mean(axis=2) / np.pi + detector_theta = d["detector_orientation_at_t"] / np.pi + # distancesb=torch.sqrt(torch.pow(diffs, 2).sum(axis=2,keepdim=True)) + distances = d["source_distance_at_t_normalized"].mean(axis=2, keepdims=True) + space_diffs = ( + d["detector_position_at_t_normalized"][:, :-1] + - d["detector_position_at_t_normalized"][:, 1:] + ) + space_delta = torch.cat( + [ + torch.zeros(b, 1, 2), + space_diffs, + ], + axis=1, + ) + + space_theta = torch.cat( + [ + torch.zeros(b, 1, 1), + (torch.atan2(space_diffs[..., 1], space_diffs[..., 0]))[:, :, None] / np.pi, + ], + axis=1, + ) + + space_dist = torch.cat( + [ + torch.zeros(b, 1, 1), + torch.sqrt(torch.pow(space_diffs, 2).sum(axis=2, keepdim=True)), + ], + axis=1, + ) + + # create the labels + labels = torch.cat( + [ + source_position, # zero center the positions + (source_theta + detector_theta + 1) % 2.0 + - 1, # initialy in units of np.pi? + distances, # try to zero center? + space_delta, + space_theta, + space_dist, + source_velocity, + ], + axis=2, + ).float() # .to(device) + # breakpoint() + # create the features + inputs = { + "drone_state": torch.cat( + [ + d["detector_position_at_t_normalized"], + times, # -times.max(axis=2,keepdim=True)[0], + space_delta, + space_theta, + space_dist, + detector_theta, + ], + dim=2, + ).float(), + "radio_feature": torch.cat( + [ + torch.log(d["beam_former_outputs_at_t"].mean(axis=2, keepdim=True)) + / 20, + d["beam_former_outputs_at_t"] + / d["beam_former_outputs_at_t"].mean( + axis=2, keepdim=True + ), # maybe pass in log values? + ], + dim=2, + ).float(), } - if 'radio_image_at_t' in d: - radio_images=d['radio_image_at_t'].float() - label_images=d['source_image_at_t'].float() - return {'inputs':inputs, - 'input_images':radio_images, - 'labels':labels, - 'label_images':label_images} - return {'inputs':inputs, - 'labels':labels} + if "radio_image_at_t" in d: + radio_images = d["radio_image_at_t"].float() + label_images = d["source_image_at_t"].float() + return { + "inputs": inputs, + "input_images": radio_images, + "labels": labels, + "label_images": label_images, + } + return {"inputs": inputs, "labels": labels} diff --git a/software/model_training_and_inference/utils/spf_generate.py b/software/model_training_and_inference/utils/spf_generate.py index daea707f..3f0e8090 100644 --- a/software/model_training_and_inference/utils/spf_generate.py +++ b/software/model_training_and_inference/utils/spf_generate.py @@ -1,4 +1,3 @@ - import argparse import os import pickle @@ -8,246 +7,298 @@ from joblib import Parallel, delayed from tqdm import tqdm from compress_pickle import dump, load -from utils.rf import NoiseWrapper, IQSource, UCADetector, ULADetector, beamformer_numba, beamformer, beamformer_old - -c=3e8 # speed of light - - -def _arctan2(x,y): - return np.arctan2(x,y) # rotation to right from x=0,y+ - -class BoundedPoint: - def __init__(self, - pos=np.ones(2)*0.5, - v=np.zeros(2), - delta_time=0.05, - width=128): - self.pos=pos - self.v=v - self.delta_time=delta_time - self.width=width - - def time_step(self,eps=1e-6): - if np.linalg.norm(self.v)==0: - return np.array(self.pos),0.0,np.zeros(2) - self.pos+=self.v*self.delta_time +from utils.rf import ( + NoiseWrapper, + IQSource, + UCADetector, + ULADetector, + beamformer_numba, + beamformer, + beamformer_old, +) - for idx in [0,1]: - while self.pos[idx]>(self.width-eps) or self.pos[idx]<0: - self.pos=np.clip(self.pos,0,self.width-eps) - self.v[idx]=-self.v[idx] - if np.linalg.norm(self.v)>0: - return np.array(self.pos),_arctan2(self.v[0],self.v[1]),self.v - return np.array(self.pos),0.0,np.zeros(2) +c = 3e8 # speed of light -def time_to_detector_offset(t,orbital_width,orbital_height,orbital_frequency=1/100.0,phase_offset=0): #1/2.0): - x=np.cos( 2*np.pi*orbital_frequency*t+phase_offset)*orbital_width - y=np.sin( 2*np.pi*orbital_frequency*t+phase_offset)*orbital_height - return np.array([ - 1/2+x, - 1/2+y]) +def _arctan2(x, y): + return np.arctan2(x, y) # rotation to right from x=0,y+ +class BoundedPoint: + def __init__(self, pos=np.ones(2) * 0.5, v=np.zeros(2), delta_time=0.05, width=128): + self.pos = pos + self.v = v + self.delta_time = delta_time + self.width = width + + def time_step(self, eps=1e-6): + if np.linalg.norm(self.v) == 0: + return np.array(self.pos), 0.0, np.zeros(2) + self.pos += self.v * self.delta_time + + for idx in [0, 1]: + while self.pos[idx] > (self.width - eps) or self.pos[idx] < 0: + self.pos = np.clip(self.pos, 0, self.width - eps) + self.v[idx] = -self.v[idx] + if np.linalg.norm(self.v) > 0: + return np.array(self.pos), _arctan2(self.v[0], self.v[1]), self.v + return np.array(self.pos), 0.0, np.zeros(2) + + +def time_to_detector_offset( + t, orbital_width, orbital_height, orbital_frequency=1 / 100.0, phase_offset=0 +): # 1/2.0): + x = np.cos(2 * np.pi * orbital_frequency * t + phase_offset) * orbital_width + y = np.sin(2 * np.pi * orbital_frequency * t + phase_offset) * orbital_height + return np.array([1 / 2 + x, 1 / 2 + y]) def generate_session(args_and_session_idx): - args,session_idx=args_and_session_idx - np.random.seed(seed=args.seed+session_idx) - wavelength=c/args.carrier_frequency - - n_sources=args.sources - n_sources_used=args.sources - if args.sources<0: - n_sources_used=np.random.choice(np.arange(1,(-args.sources)+1)) - n_sources=-args.sources - detector_speed=args.detector_speed - if args.detector_speed<0: - detector_speed=np.random.uniform(low=0.0, high=-args.detector_speed) - sigma=args.sigma - if args.sigma<0: - sigma=np.random.uniform(low=0.0, high=-args.sigma) - detector_noise=args.detector_noise - if args.detector_noise<0: - detector_noise=np.random.uniform(low=0.0, high=-args.detector_noise) - - beamformer_f=beamformer_numba if args.numba else beamformer - - if args.array_type=='linear': - d=ULADetector(args.sampling_frequency,args.elements,wavelength/2,sigma=detector_noise) # 10Mhz sampling - elif args.array_type=='circular': - d=UCADetector(args.sampling_frequency,args.elements,wavelength/4,sigma=detector_noise) # 10Mhz sampling - else: - print("Array type must be linear or circular") - sys.exit(1) - - current_source_positions=np.random.uniform(low=0, high=args.width,size=(n_sources,2)) - current_source_velocities=np.zeros((n_sources,2)) - - if args.reference: - current_source_positions=current_source_positions*0+np.array((args.width//2,args.width//4)) - sigma=0 - n_sources=1 - n_sources_used=1 - - - detector_position_phase_offset=np.random.uniform(0,2*np.pi) - detector_position_phase_offsets_at_t=np.ones((args.time_steps,1))*detector_position_phase_offset - - d.position_offset=0 - - # lets run this thing - source_positions_at_t=np.zeros((args.time_steps,n_sources,2)) - source_velocities_at_t=np.zeros((args.time_steps,n_sources,2)) - broadcasting_positions_at_t=np.zeros((args.time_steps,n_sources,1)) - broadcasting_heading_at_t=np.zeros((args.time_steps,n_sources,1)) - receiver_positions_at_t=np.zeros((args.time_steps,args.elements,2)) - source_theta_at_t=np.zeros((args.time_steps,1,1)) - source_distance_at_t=np.zeros((args.time_steps,1,1)) - - detector_orientation_at_t=np.ones((args.time_steps,1)) - - signal_matrixs_at_t=np.zeros((args.time_steps,args.elements,args.samples_per_snapshot),dtype=np.complex64) - beam_former_outputs_at_t=np.zeros((args.time_steps,args.beam_former_spacing)) - detector_position_at_t=np.zeros((args.time_steps,2)) - - thetas_at_t=np.zeros((args.time_steps,args.beam_former_spacing)) - - detector_theta=np.random.uniform(-np.pi,np.pi) - if args.reference: - detector_theta=np.random.choice([0,np.pi/4,np.pi*3/4,np.pi/2,np.pi]) - #detector_theta=np.random.choice([0,np.pi/4,np.pi/2,np.pi]) - - detector_v=np.array([np.cos(detector_theta),np.sin(detector_theta)])*detector_speed # 10m/s - - detector_initial_position=pos=np.random.uniform(0+10,args.width-10,2) - if args.fixed_detector is not None: - detector_initial_position[:2]=args.fixed_detector - - detector_bounded_point=BoundedPoint(pos=detector_initial_position, - v=detector_v, - delta_time=args.time_interval) - - source_bounded_points=[] - for idx in range(n_sources): - source_theta=np.random.uniform(-np.pi,np.pi) - source_speed=args.source_speed - if source_speed<0: - source_speed=np.random.uniform(low=0.0, high=-source_speed) - source_v=np.array( - [ - np.cos(source_theta), - np.sin(source_theta)])*source_speed - source_bounded_points.append( - BoundedPoint( - pos=np.random.uniform(0+10,args.width-10,2), - v=source_v, - delta_time=args.time_interval) - ) - - whos_broadcasting_at_t=np.random.randint(0,n_sources_used,args.time_steps) - - if args.random_emitter_timing: - emitter_p=np.random.randint(1,10,n_sources_used) - emitter_p=emitter_p/emitter_p.sum() - whos_broadcasting_at_t=np.random.choice(np.arange(n_sources_used),args.time_steps,p=emitter_p) - - if args.random_silence: - silence_p=np.random.uniform(0.0,0.8) - whos_broadcasting_at_t[np.random.choice(np.arange(args.time_steps),int(silence_p*args.time_steps),replace=False)]=-1 - - time_steps_that_broadcast=np.where(whos_broadcasting_at_t>=0)[0] - broadcasting_positions_at_t[time_steps_that_broadcast,whos_broadcasting_at_t[time_steps_that_broadcast]]=1 - #broadcasting_positions_at_t[np.arange(args.time_steps),whos_broadcasting_at_t]=1 - - #deal with source positions - #broadcasting_source_positions=source_positions_at_t[np.where(broadcasting_positions_at_t==1)[:-1]] - - #deal with detector position features - #diffs=source_positions-d['detector_position_at_t_normalized'] - #source_theta=(torch.atan2(diffs[...,1],diffs[...,0]))[:,:,None] - - - time_stamps=(np.arange(args.time_steps)*args.time_interval).reshape(-1,1) - for t_idx in np.arange(args.time_steps): - #update source positions - for idx in range(len(source_bounded_points)): - current_source_positions[idx],_,current_source_velocities[idx]=source_bounded_points[idx].time_step() - #only one source transmits at a time, TDM this part - tdm_source_idx=whos_broadcasting_at_t[t_idx] - d.rm_sources() - if tdm_source_idx>=0: - d.add_source(NoiseWrapper( - IQSource( - current_source_positions[[tdm_source_idx]], # x, y position - args.carrier_frequency), - sigma=sigma)) - - #set the detector position (its moving) - if args.detector_trajectory=='orbit': - print("There is an error in angles somewhere here") - sys.exit(1) - d.position_offset=(time_to_detector_offset(t=time_stamps[t_idx,0], - orbital_width=1/4, - orbital_height=1/3, - phase_offset=detector_position_phase_offset, - orbital_frequency=(2/3)*args.width*np.pi/detector_speed)*args.width).astype(int) - elif args.detector_trajectory=='bounce': - d.position_offset,d.orientation,_=detector_bounded_point.time_step() - detector_orientation_at_t[t_idx]=d.orientation - #if t_idx<16: - # print("WTF",d.orientation,session_idx,t_idx) - # _v=detector_bounded_point.v - # print("VEL",_v,np.arctan2(_v[1],_v[0])) - - detector_position_phase_offsets_at_t[t_idx]=detector_position_phase_offset - source_positions_at_t[t_idx]=current_source_positions - source_velocities_at_t[t_idx]=current_source_velocities - receiver_positions_at_t[t_idx]=d.all_receiver_pos() - #if t_idx<16: - # print("RECE",d.receiver_positions) - # print("ALL RECE",receiver_positions_at_t[t_idx]) - - signal_matrixs_at_t[t_idx]=d.get_signal_matrix( - start_time=time_stamps[t_idx,0], - duration=args.samples_per_snapshot/d.sampling_frequency) - thetas_at_t[t_idx],beam_former_outputs_at_t[t_idx],_=beamformer_f( - d.all_receiver_pos(with_offset=False), - signal_matrixs_at_t[t_idx], - args.carrier_frequency,spacing=args.beam_former_spacing, - offset=d.orientation) - #print(d.orientation,detector_theta) - detector_position_at_t[t_idx]=d.position_offset - - if tdm_source_idx>=0: - diff=current_source_positions[tdm_source_idx]-detector_position_at_t[t_idx] - #both of these are in regular units, radians to the right of x+ - #but it doesnt matter because we just want the difference - source_theta_at_t[t_idx]=(_arctan2(diff[[0]],diff[[1]])-d.orientation+np.pi)%(2*np.pi)-np.pi - source_distance_at_t[t_idx]=np.sqrt(np.power(diff,2).sum()) + args, session_idx = args_and_session_idx + np.random.seed(seed=args.seed + session_idx) + wavelength = c / args.carrier_frequency + + n_sources = args.sources + n_sources_used = args.sources + if args.sources < 0: + n_sources_used = np.random.choice(np.arange(1, (-args.sources) + 1)) + n_sources = -args.sources + detector_speed = args.detector_speed + if args.detector_speed < 0: + detector_speed = np.random.uniform(low=0.0, high=-args.detector_speed) + sigma = args.sigma + if args.sigma < 0: + sigma = np.random.uniform(low=0.0, high=-args.sigma) + detector_noise = args.detector_noise + if args.detector_noise < 0: + detector_noise = np.random.uniform(low=0.0, high=-args.detector_noise) + + beamformer_f = beamformer_numba if args.numba else beamformer + + if args.array_type == "linear": + d = ULADetector( + args.sampling_frequency, args.elements, wavelength / 2, sigma=detector_noise + ) # 10Mhz sampling + elif args.array_type == "circular": + d = UCADetector( + args.sampling_frequency, args.elements, wavelength / 4, sigma=detector_noise + ) # 10Mhz sampling else: - source_theta_at_t[t_idx]=0 #(np.arctan2(diff[[1]],diff[[0]])-d.orientation+np.pi)%(2*np.pi)-np.pi - source_distance_at_t[t_idx]=0 #np.sqrt(np.power(diff,2).sum()) - session={ - 'broadcasting_positions_at_t':broadcasting_positions_at_t, # list of (time_steps,sources,1) - 'source_positions_at_t':source_positions_at_t, # (time_steps,sources,2[x,y]) - 'source_velocities_at_t':source_velocities_at_t, # (time_steps,sources,2[x,y]) - 'receiver_positions_at_t':receiver_positions_at_t, # (time_steps,receivers,2[x,y]) - 'signal_matrixs_at_t':signal_matrixs_at_t, # (time_steps,receivers,samples_per_snapshot) - 'beam_former_outputs_at_t':beam_former_outputs_at_t, #(timesteps,thetas_tested_for_steering) - 'thetas_at_t':thetas_at_t, #(timesteps,thetas_tested_for_steering) - 'detector_position_phase_offsets_at_t':detector_position_phase_offsets_at_t, # (timesteps,1) # phase offset for orbit dynamics - 'time_stamps':time_stamps, # (timestamps,1) # the time for each step - 'width_at_t':np.ones((args.time_steps,1),dtype=int)*args.width, # (timestamps, 1) # width in meters - 'detector_orientation_at_t':detector_orientation_at_t, # (timesteps,1 ) # orientation of the detector at t - 'detector_position_at_t':detector_position_at_t, # (time_steps,2[x,y]) - 'source_theta_at_t':source_theta_at_t, # (time_steps, 1) # if a source is broadcasting than the orientation from detector oreintation to source, else 0 - 'source_distance_at_t':source_distance_at_t, # (time_steps,1) # if a source is broadcasting than this is the distance to that source from the detector - } - return session + print("Array type must be linear or circular") + sys.exit(1) + + current_source_positions = np.random.uniform( + low=0, high=args.width, size=(n_sources, 2) + ) + current_source_velocities = np.zeros((n_sources, 2)) + + if args.reference: + current_source_positions = current_source_positions * 0 + np.array( + (args.width // 2, args.width // 4) + ) + sigma = 0 + n_sources = 1 + n_sources_used = 1 + + detector_position_phase_offset = np.random.uniform(0, 2 * np.pi) + detector_position_phase_offsets_at_t = ( + np.ones((args.time_steps, 1)) * detector_position_phase_offset + ) + + d.position_offset = 0 + + # lets run this thing + source_positions_at_t = np.zeros((args.time_steps, n_sources, 2)) + source_velocities_at_t = np.zeros((args.time_steps, n_sources, 2)) + broadcasting_positions_at_t = np.zeros((args.time_steps, n_sources, 1)) + broadcasting_heading_at_t = np.zeros((args.time_steps, n_sources, 1)) + receiver_positions_at_t = np.zeros((args.time_steps, args.elements, 2)) + source_theta_at_t = np.zeros((args.time_steps, 1, 1)) + source_distance_at_t = np.zeros((args.time_steps, 1, 1)) + + detector_orientation_at_t = np.ones((args.time_steps, 1)) + + signal_matrixs_at_t = np.zeros( + (args.time_steps, args.elements, args.samples_per_snapshot), dtype=np.complex64 + ) + beam_former_outputs_at_t = np.zeros((args.time_steps, args.beam_former_spacing)) + detector_position_at_t = np.zeros((args.time_steps, 2)) + + thetas_at_t = np.zeros((args.time_steps, args.beam_former_spacing)) + + detector_theta = np.random.uniform(-np.pi, np.pi) + if args.reference: + detector_theta = np.random.choice( + [0, np.pi / 4, np.pi * 3 / 4, np.pi / 2, np.pi] + ) + # detector_theta=np.random.choice([0,np.pi/4,np.pi/2,np.pi]) + + detector_v = ( + np.array([np.cos(detector_theta), np.sin(detector_theta)]) * detector_speed + ) # 10m/s + + detector_initial_position = pos = np.random.uniform(0 + 10, args.width - 10, 2) + if args.fixed_detector is not None: + detector_initial_position[:2] = args.fixed_detector + + detector_bounded_point = BoundedPoint( + pos=detector_initial_position, v=detector_v, delta_time=args.time_interval + ) + + source_bounded_points = [] + for idx in range(n_sources): + source_theta = np.random.uniform(-np.pi, np.pi) + source_speed = args.source_speed + if source_speed < 0: + source_speed = np.random.uniform(low=0.0, high=-source_speed) + source_v = np.array([np.cos(source_theta), np.sin(source_theta)]) * source_speed + source_bounded_points.append( + BoundedPoint( + pos=np.random.uniform(0 + 10, args.width - 10, 2), + v=source_v, + delta_time=args.time_interval, + ) + ) + + whos_broadcasting_at_t = np.random.randint(0, n_sources_used, args.time_steps) + + if args.random_emitter_timing: + emitter_p = np.random.randint(1, 10, n_sources_used) + emitter_p = emitter_p / emitter_p.sum() + whos_broadcasting_at_t = np.random.choice( + np.arange(n_sources_used), args.time_steps, p=emitter_p + ) + + if args.random_silence: + silence_p = np.random.uniform(0.0, 0.8) + whos_broadcasting_at_t[ + np.random.choice( + np.arange(args.time_steps), + int(silence_p * args.time_steps), + replace=False, + ) + ] = -1 + + time_steps_that_broadcast = np.where(whos_broadcasting_at_t >= 0)[0] + broadcasting_positions_at_t[ + time_steps_that_broadcast, whos_broadcasting_at_t[time_steps_that_broadcast] + ] = 1 + # broadcasting_positions_at_t[np.arange(args.time_steps),whos_broadcasting_at_t]=1 + + # deal with source positions + # broadcasting_source_positions=source_positions_at_t[np.where(broadcasting_positions_at_t==1)[:-1]] + + # deal with detector position features + # diffs=source_positions-d['detector_position_at_t_normalized'] + # source_theta=(torch.atan2(diffs[...,1],diffs[...,0]))[:,:,None] + + time_stamps = (np.arange(args.time_steps) * args.time_interval).reshape(-1, 1) + for t_idx in np.arange(args.time_steps): + # update source positions + for idx in range(len(source_bounded_points)): + ( + current_source_positions[idx], + _, + current_source_velocities[idx], + ) = source_bounded_points[idx].time_step() + # only one source transmits at a time, TDM this part + tdm_source_idx = whos_broadcasting_at_t[t_idx] + d.rm_sources() + if tdm_source_idx >= 0: + d.add_source( + NoiseWrapper( + IQSource( + current_source_positions[[tdm_source_idx]], # x, y position + args.carrier_frequency, + ), + sigma=sigma, + ) + ) + + # set the detector position (its moving) + if args.detector_trajectory == "orbit": + print("There is an error in angles somewhere here") + sys.exit(1) + d.position_offset = ( + time_to_detector_offset( + t=time_stamps[t_idx, 0], + orbital_width=1 / 4, + orbital_height=1 / 3, + phase_offset=detector_position_phase_offset, + orbital_frequency=(2 / 3) * args.width * np.pi / detector_speed, + ) + * args.width + ).astype(int) + elif args.detector_trajectory == "bounce": + d.position_offset, d.orientation, _ = detector_bounded_point.time_step() + detector_orientation_at_t[t_idx] = d.orientation + # if t_idx<16: + # print("WTF",d.orientation,session_idx,t_idx) + # _v=detector_bounded_point.v + # print("VEL",_v,np.arctan2(_v[1],_v[0])) + + detector_position_phase_offsets_at_t[t_idx] = detector_position_phase_offset + source_positions_at_t[t_idx] = current_source_positions + source_velocities_at_t[t_idx] = current_source_velocities + receiver_positions_at_t[t_idx] = d.all_receiver_pos() + # if t_idx<16: + # print("RECE",d.receiver_positions) + # print("ALL RECE",receiver_positions_at_t[t_idx]) + + signal_matrixs_at_t[t_idx] = d.get_signal_matrix( + start_time=time_stamps[t_idx, 0], + duration=args.samples_per_snapshot / d.sampling_frequency, + ) + thetas_at_t[t_idx], beam_former_outputs_at_t[t_idx], _ = beamformer_f( + d.all_receiver_pos(with_offset=False), + signal_matrixs_at_t[t_idx], + args.carrier_frequency, + spacing=args.beam_former_spacing, + offset=d.orientation, + ) + # print(d.orientation,detector_theta) + detector_position_at_t[t_idx] = d.position_offset + + if tdm_source_idx >= 0: + diff = ( + current_source_positions[tdm_source_idx] - detector_position_at_t[t_idx] + ) + # both of these are in regular units, radians to the right of x+ + # but it doesnt matter because we just want the difference + source_theta_at_t[t_idx] = ( + _arctan2(diff[[0]], diff[[1]]) - d.orientation + np.pi + ) % (2 * np.pi) - np.pi + source_distance_at_t[t_idx] = np.sqrt(np.power(diff, 2).sum()) + else: + source_theta_at_t[ + t_idx + ] = 0 # (np.arctan2(diff[[1]],diff[[0]])-d.orientation+np.pi)%(2*np.pi)-np.pi + source_distance_at_t[t_idx] = 0 # np.sqrt(np.power(diff,2).sum()) + session = { + "broadcasting_positions_at_t": broadcasting_positions_at_t, # list of (time_steps,sources,1) + "source_positions_at_t": source_positions_at_t, # (time_steps,sources,2[x,y]) + "source_velocities_at_t": source_velocities_at_t, # (time_steps,sources,2[x,y]) + "receiver_positions_at_t": receiver_positions_at_t, # (time_steps,receivers,2[x,y]) + "signal_matrixs_at_t": signal_matrixs_at_t, # (time_steps,receivers,samples_per_snapshot) + "beam_former_outputs_at_t": beam_former_outputs_at_t, # (timesteps,thetas_tested_for_steering) + "thetas_at_t": thetas_at_t, # (timesteps,thetas_tested_for_steering) + "detector_position_phase_offsets_at_t": detector_position_phase_offsets_at_t, # (timesteps,1) # phase offset for orbit dynamics + "time_stamps": time_stamps, # (timestamps,1) # the time for each step + "width_at_t": np.ones((args.time_steps, 1), dtype=int) + * args.width, # (timestamps, 1) # width in meters + "detector_orientation_at_t": detector_orientation_at_t, # (timesteps,1 ) # orientation of the detector at t + "detector_position_at_t": detector_position_at_t, # (time_steps,2[x,y]) + "source_theta_at_t": source_theta_at_t, # (time_steps, 1) # if a source is broadcasting than the orientation from detector oreintation to source, else 0 + "source_distance_at_t": source_distance_at_t, # (time_steps,1) # if a source is broadcasting than this is the distance to that source from the detector + } + return session -def generate_session_and_dump(args_and_session_idx): - args,session_idx=args_and_session_idx - session=generate_session(args_and_session_idx) - dump(session,"/".join([args.output,'session_%08d.pkl' % session_idx]),compression="lzma") +def generate_session_and_dump(args_and_session_idx): + args, session_idx = args_and_session_idx + session = generate_session(args_and_session_idx) + dump( + session, + "/".join([args.output, "session_%08d.pkl" % session_idx]), + compression="lzma", + ) diff --git a/software/sdrpluto/01_phase_sync.py b/software/sdrpluto/01_phase_sync.py index 2a413f96..4e0fb2de 100644 --- a/software/sdrpluto/01_phase_sync.py +++ b/software/sdrpluto/01_phase_sync.py @@ -1,107 +1,122 @@ -#Starter code from jon Kraft +# Starter code from jon Kraft import adi -#;import matplotlib.pyplot as plt + +# ;import matplotlib.pyplot as plt import numpy as np from math import lcm from sdr import * -c=3e8 +c = 3e8 fc0 = int(200e3) -fs = int(12e6) # must be <=30.72 MHz if both channels are enabled -#rx_lo = int(2.4e9) #4e9 -rx_lo = int(2.5e9) #4e9 +fs = int(12e6) # must be <=30.72 MHz if both channels are enabled +# rx_lo = int(2.4e9) #4e9 +rx_lo = int(2.5e9) # 4e9 tx_lo = rx_lo rx_mode = "manual" # can be "manual" or "slow_attack" -rx_gain = -30 +rx_gain = -30 tx_gain = -30 -wavelength = c/rx_lo # wavelength of the RF carrier +wavelength = c / rx_lo # wavelength of the RF carrier -rx_n=int(2**10) +rx_n = int(2**10) -sdr = adi.ad9361(uri='ip:192.168.3.1') +sdr = adi.ad9361(uri="ip:192.168.3.1") -#setup RX +# setup RX sdr.rx_enabled_channels = [0, 1] sdr.sample_rate = fs -assert(sdr.sample_rate==fs) -sdr.rx_rf_bandwidth = int(fc0*3) +assert sdr.sample_rate == fs +sdr.rx_rf_bandwidth = int(fc0 * 3) sdr.rx_lo = int(rx_lo) sdr.gain_control_mode = rx_mode sdr.rx_hardwaregain_chan0 = int(rx_gain) sdr.rx_hardwaregain_chan1 = int(rx_gain) sdr.rx_buffer_size = int(rx_n) -sdr._rxadc.set_kernel_buffers_count(1) # set buffers to 1 (instead of the default 4) to avoid stale data on Pluto +sdr._rxadc.set_kernel_buffers_count( + 1 +) # set buffers to 1 (instead of the default 4) to avoid stale data on Pluto -#setup TX -sdr.tx_enabled_channels = [0,1] +# setup TX +sdr.tx_enabled_channels = [0, 1] sdr.tx_rf_bandwidth = int(fc0) sdr.tx_lo = int(tx_lo) -sdr.tx_cyclic_buffer = True # this keeps repeating! -sdr.tx_hardwaregain_chan0 = int(-88) #tx_gain) -sdr.tx_hardwaregain_chan1 = int(tx_gain) # use Tx2 for calibration -tx_n=int(min(lcm(fc0,fs),rx_n*8)) #1024*1024*1024) # tx for longer than rx -sdr.tx_buffer_size = tx_n*2 #tx_n - -#since its a cyclic buffer its important to end on a full phase -t = np.arange(0, tx_n)/fs -iq0 = np.exp(1j*2*np.pi*t*fc0)*(2**14) -sdr.tx([iq0,iq0]) # Send Tx data. - -detector=ULADetector(fs,2,wavelength/2) - -def sample_phase_offset_rx(iterations=64,calibration=1+0j): - steerings=[] - for _ in np.arange(iterations): - signal_matrix=np.vstack(sdr.rx()) - thetas,sds,steering=beamformer(detector,signal_matrix,rx_lo,spacing=2*1024+1,calibration=calibration) - steerings.append(steering[sds.argmax()]) - return np.array(steerings).mean(axis=0) - -calibration=np.conjugate(sample_phase_offset_rx()) -print("calibration complete",calibration) -calibration_new=np.conjugate(sample_phase_offset_rx(calibration=calibration)) -print("new cal",calibration_new) +sdr.tx_cyclic_buffer = True # this keeps repeating! +sdr.tx_hardwaregain_chan0 = int(-88) # tx_gain) +sdr.tx_hardwaregain_chan1 = int(tx_gain) # use Tx2 for calibration +tx_n = int(min(lcm(fc0, fs), rx_n * 8)) # 1024*1024*1024) # tx for longer than rx +sdr.tx_buffer_size = tx_n * 2 # tx_n + +# since its a cyclic buffer its important to end on a full phase +t = np.arange(0, tx_n) / fs +iq0 = np.exp(1j * 2 * np.pi * t * fc0) * (2**14) +sdr.tx([iq0, iq0]) # Send Tx data. + +detector = ULADetector(fs, 2, wavelength / 2) + + +def sample_phase_offset_rx(iterations=64, calibration=1 + 0j): + steerings = [] + for _ in np.arange(iterations): + signal_matrix = np.vstack(sdr.rx()) + thetas, sds, steering = beamformer( + detector, + signal_matrix, + rx_lo, + spacing=2 * 1024 + 1, + calibration=calibration, + ) + steerings.append(steering[sds.argmax()]) + return np.array(steerings).mean(axis=0) + + +calibration = np.conjugate(sample_phase_offset_rx()) +print("calibration complete", calibration) +calibration_new = np.conjugate(sample_phase_offset_rx(calibration=calibration)) +print("new cal", calibration_new) sdr.tx_destroy_buffer() sdr.rx_destroy_buffer() -#setup RX +# setup RX sdr.rx_enabled_channels = [0, 1] sdr.sample_rate = fs -assert(sdr.sample_rate==fs) +assert sdr.sample_rate == fs sdr.rx_rf_bandwidth = int(fc0) sdr.rx_lo = int(rx_lo) sdr.gain_control_mode = rx_mode sdr.rx_hardwaregain_chan0 = int(rx_gain) sdr.rx_hardwaregain_chan1 = int(rx_gain) sdr.rx_buffer_size = int(rx_n) -sdr._rxadc.set_kernel_buffers_count(1) # set buffers to 1 (instead of the default 4) to avoid stale data on Pluto +sdr._rxadc.set_kernel_buffers_count( + 1 +) # set buffers to 1 (instead of the default 4) to avoid stale data on Pluto -#setup TX +# setup TX sdr.tx_enabled_channels = [0] sdr.tx_rf_bandwidth = int(fc0) sdr.tx_lo = int(tx_lo) -sdr.tx_cyclic_buffer = True # this keeps repeating! -sdr.tx_hardwaregain_chan0 = int(tx_gain) #tx_gain) -sdr.tx_hardwaregain_chan1 = int(-80) # use Tx2 for calibration +sdr.tx_cyclic_buffer = True # this keeps repeating! +sdr.tx_hardwaregain_chan0 = int(tx_gain) # tx_gain) +sdr.tx_hardwaregain_chan1 = int(-80) # use Tx2 for calibration -tx_n=int(min(lcm(fc0,fs),rx_n*8)) #1024*1024*1024) # tx for longer than rx +tx_n = int(min(lcm(fc0, fs), rx_n * 8)) # 1024*1024*1024) # tx for longer than rx sdr.tx_buffer_size = tx_n -#since its a cyclic buffer its important to end on a full phase -t = np.arange(0, tx_n)/fs -iq0 = np.exp(1j*2*np.pi*t*fc0)*(2**14) +# since its a cyclic buffer its important to end on a full phase +t = np.arange(0, tx_n) / fs +iq0 = np.exp(1j * 2 * np.pi * t * fc0) * (2**14) sdr.tx(iq0) # Send Tx data. -fig,axs=plt.subplots(2,1,figsize=(4,8)) +fig, axs = plt.subplots(2, 1, figsize=(4, 8)) while True: - axs[0].cla() - signal_matrix=np.vstack(sdr.rx()) - thetas,sds,steering=beamformer(detector,signal_matrix,rx_lo,spacing=2*512+1,calibration=calibration) - axs[0].scatter(360*thetas/(2*np.pi),sds,s=0.5) - plt.draw() - plt.pause(0.01) + axs[0].cla() + signal_matrix = np.vstack(sdr.rx()) + thetas, sds, steering = beamformer( + detector, signal_matrix, rx_lo, spacing=2 * 512 + 1, calibration=calibration + ) + axs[0].scatter(360 * thetas / (2 * np.pi), sds, s=0.5) + plt.draw() + plt.pause(0.01) diff --git a/software/sdrpluto/02_wifi_direction.py b/software/sdrpluto/02_wifi_direction.py index c72b4428..4069f170 100644 --- a/software/sdrpluto/02_wifi_direction.py +++ b/software/sdrpluto/02_wifi_direction.py @@ -1,113 +1,130 @@ -#Starter code from jon Kraft +# Starter code from jon Kraft import adi -#;import matplotlib.pyplot as plt + +# ;import matplotlib.pyplot as plt import numpy as np from math import lcm -#from sdr import * + +# from sdr import * -c=3e8 +c = 3e8 fc0 = int(500e3) -fs = int(4e6) # must be <=30.72 MHz if both channels are enabled -rx_lo = int(2.45e9) #4e9 +fs = int(4e6) # must be <=30.72 MHz if both channels are enabled +rx_lo = int(2.45e9) # 4e9 tx_lo = rx_lo rx_mode = "manual" # can be "manual" or "slow_attack" -rx_gain = -30 +rx_gain = -30 tx_gain = -30 -wavelength = c/rx_lo # wavelength of the RF carrier +wavelength = c / rx_lo # wavelength of the RF carrier -rx_n=int(2**10) +rx_n = int(2**10) -sdr = adi.ad9361(uri='ip:192.168.3.1') +sdr = adi.ad9361(uri="ip:192.168.3.1") -#setup RX +# setup RX sdr.rx_enabled_channels = [0, 1] sdr.sample_rate = fs -assert(sdr.sample_rate==fs) -sdr.rx_rf_bandwidth = int(fc0*3) +assert sdr.sample_rate == fs +sdr.rx_rf_bandwidth = int(fc0 * 3) sdr.rx_lo = int(rx_lo) sdr.gain_control_mode = rx_mode sdr.rx_hardwaregain_chan0 = int(rx_gain) sdr.rx_hardwaregain_chan1 = int(rx_gain) sdr.rx_buffer_size = int(rx_n) -sdr._rxadc.set_kernel_buffers_count(1) # set buffers to 1 (instead of the default 4) to avoid stale data on Pluto +sdr._rxadc.set_kernel_buffers_count( + 1 +) # set buffers to 1 (instead of the default 4) to avoid stale data on Pluto -#setup TX -sdr.tx_enabled_channels = [0,1] -sdr.tx_rf_bandwidth = int(fc0*3) +# setup TX +sdr.tx_enabled_channels = [0, 1] +sdr.tx_rf_bandwidth = int(fc0 * 3) sdr.tx_lo = int(tx_lo) -sdr.tx_cyclic_buffer = True # this keeps repeating! -sdr.tx_hardwaregain_chan0 = int(-88) #tx_gain) -sdr.tx_hardwaregain_chan1 = int(tx_gain) # use Tx2 for calibration -tx_n=int(min(lcm(fc0,fs),rx_n*8)) #1024*1024*1024) # tx for longer than rx -sdr.tx_buffer_size = tx_n*2 #tx_n - -#since its a cyclic buffer its important to end on a full phase -t = np.arange(0, tx_n)/fs -iq0 = np.exp(1j*2*np.pi*t*fc0)*(2**14) -sdr.tx([iq0,iq0]) # Send Tx data. - -detector=ULADetector(fs,2,wavelength/2) - -def sample_phase_offset_rx(iterations=64,calibration=1+0j): - steerings=[] - for _ in np.arange(iterations): - signal_matrix=np.vstack(sdr.rx()) - thetas,sds,steering=beamformer(detector,signal_matrix,rx_lo,spacing=2*1024+1,calibration=calibration) - steerings.append(steering[sds.argmax()]) - return np.array(steerings).mean(axis=0) - -calibration=np.conjugate(sample_phase_offset_rx()) -print("calibration complete",calibration) -calibration_new=np.conjugate(sample_phase_offset_rx(calibration=calibration)) -print("new cal",calibration_new) +sdr.tx_cyclic_buffer = True # this keeps repeating! +sdr.tx_hardwaregain_chan0 = int(-88) # tx_gain) +sdr.tx_hardwaregain_chan1 = int(tx_gain) # use Tx2 for calibration +tx_n = int(min(lcm(fc0, fs), rx_n * 8)) # 1024*1024*1024) # tx for longer than rx +sdr.tx_buffer_size = tx_n * 2 # tx_n + +# since its a cyclic buffer its important to end on a full phase +t = np.arange(0, tx_n) / fs +iq0 = np.exp(1j * 2 * np.pi * t * fc0) * (2**14) +sdr.tx([iq0, iq0]) # Send Tx data. + +detector = ULADetector(fs, 2, wavelength / 2) + + +def sample_phase_offset_rx(iterations=64, calibration=1 + 0j): + steerings = [] + for _ in np.arange(iterations): + signal_matrix = np.vstack(sdr.rx()) + thetas, sds, steering = beamformer( + detector, + signal_matrix, + rx_lo, + spacing=2 * 1024 + 1, + calibration=calibration, + ) + steerings.append(steering[sds.argmax()]) + return np.array(steerings).mean(axis=0) + + +calibration = np.conjugate(sample_phase_offset_rx()) +print("calibration complete", calibration) +calibration_new = np.conjugate(sample_phase_offset_rx(calibration=calibration)) +print("new cal", calibration_new) sdr.tx_destroy_buffer() sdr.rx_destroy_buffer() -#setup RX +# setup RX sdr.rx_enabled_channels = [0, 1] sdr.sample_rate = fs -assert(sdr.sample_rate==fs) -sdr.rx_rf_bandwidth = int(fc0*3) +assert sdr.sample_rate == fs +sdr.rx_rf_bandwidth = int(fc0 * 3) sdr.rx_lo = int(rx_lo) sdr.gain_control_mode = rx_mode sdr.rx_hardwaregain_chan0 = int(rx_gain) sdr.rx_hardwaregain_chan1 = int(rx_gain) sdr.rx_buffer_size = int(rx_n) -sdr._rxadc.set_kernel_buffers_count(1) # set buffers to 1 (instead of the default 4) to avoid stale data on Pluto +sdr._rxadc.set_kernel_buffers_count( + 1 +) # set buffers to 1 (instead of the default 4) to avoid stale data on Pluto -#setup TX +# setup TX sdr.tx_enabled_channels = [] -sdr.tx_rf_bandwidth = int(fc0*3) +sdr.tx_rf_bandwidth = int(fc0 * 3) sdr.tx_lo = int(tx_lo) -sdr.tx_cyclic_buffer = True # this keeps repeating! -sdr.tx_hardwaregain_chan0 = int(tx_gain) #tx_gain) #tx_gain) -sdr.tx_hardwaregain_chan1 = int(-80) # use Tx2 for calibration +sdr.tx_cyclic_buffer = True # this keeps repeating! +sdr.tx_hardwaregain_chan0 = int(tx_gain) # tx_gain) #tx_gain) +sdr.tx_hardwaregain_chan1 = int(-80) # use Tx2 for calibration # -tx_n=int(min(lcm(fc0,fs),rx_n*8)) #1024*1024*1024) # tx for longer than rx +tx_n = int(min(lcm(fc0, fs), rx_n * 8)) # 1024*1024*1024) # tx for longer than rx sdr.tx_buffer_size = tx_n -#since its a cyclic buffer its important to end on a full phase -t = np.arange(0, tx_n)/fs -iq0 = np.exp(1j*2*np.pi*t*fc0)*(2**14) +# since its a cyclic buffer its important to end on a full phase +t = np.arange(0, tx_n) / fs +iq0 = np.exp(1j * 2 * np.pi * t * fc0) * (2**14) sdr.tx(iq0) # Send Tx data. import time -fig,axs=plt.subplots(1,1,figsize=(4,4)) -intervals=2*64+1 -counts=np.zeros(intervals-1) +fig, axs = plt.subplots(1, 1, figsize=(4, 4)) + +intervals = 2 * 64 + 1 +counts = np.zeros(intervals - 1) while True: - signal_matrix=np.vstack(sdr.rx()) - thetas,sds,steering=beamformer(detector,signal_matrix,rx_lo,spacing=intervals,calibration=calibration) - if sds.max()>1000: - print(time.time(),sds.max(),thetas[sds.argmax()]) - counts[sds.argmax()%(intervals-1)]+=1 - axs.cla() - axs.stairs(counts, 360*thetas/(2*np.pi)) - #axs.scatter(360*thetas/(2*np.pi),sds,s=0.5) - plt.draw() - plt.pause(0.01) + signal_matrix = np.vstack(sdr.rx()) + thetas, sds, steering = beamformer( + detector, signal_matrix, rx_lo, spacing=intervals, calibration=calibration + ) + if sds.max() > 1000: + print(time.time(), sds.max(), thetas[sds.argmax()]) + counts[sds.argmax() % (intervals - 1)] += 1 + axs.cla() + axs.stairs(counts, 360 * thetas / (2 * np.pi)) + # axs.scatter(360*thetas/(2*np.pi),sds,s=0.5) + plt.draw() + plt.pause(0.01) diff --git a/software/sdrpluto/gather.py b/software/sdrpluto/gather.py index f8005a64..bed112b2 100644 --- a/software/sdrpluto/gather.py +++ b/software/sdrpluto/gather.py @@ -5,294 +5,337 @@ from math import gcd import matplotlib.pyplot as plt import time -c=3e8 + +c = 3e8 + def setup_rxtx_and_phase_calibration(args): print("Starting inter antenna receiver phase calibration") fc0 = int(args.fi) - fs = int(args.fs) # must be <=30.72 MHz if both channels are enabled - rx_lo = int(args.fc) #4e9 + fs = int(args.fs) # must be <=30.72 MHz if both channels are enabled + rx_lo = int(args.fc) # 4e9 tx_lo = rx_lo # setup receive - rx_mode = args.rx_mode - rx_gain = args.rx_gain + rx_mode = args.rx_mode + rx_gain = args.rx_gain - rx_n=args.rx_n + rx_n = args.rx_n - tx_gain_calibration=-50 + tx_gain_calibration = -50 - emitter_online=False - retries=0 - while retries<10: - #try to setup TX and lets see if it works - #sdr_emitter = adi.ad9361(uri='ip:%s' % args.emitter_ip) - sdr_rxtx = adi.ad9361(uri='ip:%s' % args.receiver_ip) + emitter_online = False + retries = 0 + while retries < 10: + # try to setup TX and lets see if it works + # sdr_emitter = adi.ad9361(uri='ip:%s' % args.emitter_ip) + sdr_rxtx = adi.ad9361(uri="ip:%s" % args.receiver_ip) sdr_rxtx.rx_enabled_channels = [0, 1] sdr_rxtx.sample_rate = fs - assert(sdr_rxtx.sample_rate==fs) - sdr_rxtx.rx_rf_bandwidth = int(3*args.fi) #fc0*5) #TODO! + assert sdr_rxtx.sample_rate == fs + sdr_rxtx.rx_rf_bandwidth = int(3 * args.fi) # fc0*5) #TODO! sdr_rxtx.rx_lo = int(rx_lo) sdr_rxtx.gain_control_mode = rx_mode sdr_rxtx.rx_hardwaregain_chan0 = int(rx_gain) sdr_rxtx.rx_hardwaregain_chan1 = int(rx_gain) sdr_rxtx.rx_buffer_size = int(rx_n) - sdr_rxtx._rxadc.set_kernel_buffers_count(2) # set buffers to 1 (instead of the default 4) to avoid stale data on Pluto + sdr_rxtx._rxadc.set_kernel_buffers_count( + 2 + ) # set buffers to 1 (instead of the default 4) to avoid stale data on Pluto - #drop the first bunch of frames + # drop the first bunch of frames for _ in range(20): sdr_rxtx.rx() - #setup TX - sdr_rxtx.tx_rf_bandwidth = int(3*args.fi) - assert(sdr_rxtx.tx_rf_bandwidth==int(3*args.fi)) + # setup TX + sdr_rxtx.tx_rf_bandwidth = int(3 * args.fi) + assert sdr_rxtx.tx_rf_bandwidth == int(3 * args.fi) sdr_rxtx.tx_lo = int(tx_lo) - assert(sdr_rxtx.tx_lo==tx_lo) + assert sdr_rxtx.tx_lo == tx_lo sdr_rxtx.tx_enabled_channels = [1] - sdr_rxtx.tx_hardwaregain_chan0 = int(-80) #tx_gain) #tx_gain) - sdr_rxtx.tx_hardwaregain_chan1 = int(tx_gain_calibration) # use Tx2 for calibration - assert(sdr_rxtx.tx_hardwaregain_chan1==int(tx_gain_calibration)) + sdr_rxtx.tx_hardwaregain_chan0 = int(-80) # tx_gain) #tx_gain) + sdr_rxtx.tx_hardwaregain_chan1 = int( + tx_gain_calibration + ) # use Tx2 for calibration + assert sdr_rxtx.tx_hardwaregain_chan1 == int(tx_gain_calibration) # - tx_n=int(fs/gcd(fs,fc0)) - while tx_n<1024*16: - tx_n*=2 - #sdr_rxtx.tx_buffer_size = tx_n - - #since its a cyclic buffer its important to end on a full phase - t = np.arange(0, tx_n)/fs # time at each point assuming we are sending samples at (1/fs)s - iq0 = np.exp(1j*2*np.pi*fc0*t)*(2**14) - #try to reset the tx + tx_n = int(fs / gcd(fs, fc0)) + while tx_n < 1024 * 16: + tx_n *= 2 + # sdr_rxtx.tx_buffer_size = tx_n + + # since its a cyclic buffer its important to end on a full phase + t = ( + np.arange(0, tx_n) / fs + ) # time at each point assuming we are sending samples at (1/fs)s + iq0 = np.exp(1j * 2 * np.pi * fc0 * t) * (2**14) + # try to reset the tx sdr_rxtx.tx_destroy_buffer() - sdr_rxtx.tx_cyclic_buffer = True # this keeps repeating! - assert(sdr_rxtx.tx_cyclic_buffer==True) + sdr_rxtx.tx_cyclic_buffer = True # this keeps repeating! + assert sdr_rxtx.tx_cyclic_buffer == True sdr_rxtx.tx(iq0) # Send Tx data. - - #give RX a chance to calm down + + # give RX a chance to calm down for _ in range(10): sdr_rxtx.rx() - #test to see what frequency we are seeing - freq = np.fft.fftfreq(rx_n,d=1.0/fs) - signal_matrix=np.vstack(sdr_rxtx.rx()) + # test to see what frequency we are seeing + freq = np.fft.fftfreq(rx_n, d=1.0 / fs) + signal_matrix = np.vstack(sdr_rxtx.rx()) sp = np.fft.fft(signal_matrix[0]) - max_freq=freq[np.abs(np.argmax(sp.real))] - if np.abs(max_freq-args.fi)<(args.fs/rx_n+1): + max_freq = freq[np.abs(np.argmax(sp.real))] + if np.abs(max_freq - args.fi) < (args.fs / rx_n + 1): print("Emitter online after %d retries" % retries) - emitter_online=True + emitter_online = True break - retries+=1 - sdr_rxtx=None + retries += 1 + sdr_rxtx = None time.sleep(1) - if emitter_online==False: - print("Failed to bring emitter online") - return None + if emitter_online == False: + print("Failed to bring emitter online") + return None - #get some new data + # get some new data for retry in range(20): - n_calibration_frames=200 - phase_calibrations=np.zeros(n_calibration_frames) - for idx in range(n_calibration_frames): - sdr_rxtx.rx() - phase_calibrations[idx]=((np.angle(signal_matrix[0])-np.angle(signal_matrix[1]))%(2*np.pi)).mean() # TODO THIS BREAKS if diff is near 2*np.pi... - if phase_calibrations.std()<1e-5: - sdr_rxtx.tx_destroy_buffer() - print("Final phase calibration (radians) is %0.4f" % phase_calibrations.mean(),"(fraction of 2pi) %0.4f" % (phase_calibrations.mean()/(2*np.pi))) - sdr_rxtx.phase_calibration=phase_calibrations.mean() - return sdr_rxtx + n_calibration_frames = 200 + phase_calibrations = np.zeros(n_calibration_frames) + for idx in range(n_calibration_frames): + sdr_rxtx.rx() + phase_calibrations[idx] = ( + (np.angle(signal_matrix[0]) - np.angle(signal_matrix[1])) % (2 * np.pi) + ).mean() # TODO THIS BREAKS if diff is near 2*np.pi... + if phase_calibrations.std() < 1e-5: + sdr_rxtx.tx_destroy_buffer() + print( + "Final phase calibration (radians) is %0.4f" + % phase_calibrations.mean(), + "(fraction of 2pi) %0.4f" % (phase_calibrations.mean() / (2 * np.pi)), + ) + sdr_rxtx.phase_calibration = phase_calibrations.mean() + return sdr_rxtx sdr_rxtx.tx_destroy_buffer() return None - def setup_rx_and_tx(args): fc0 = int(args.fi) - fs = int(args.fs) # must be <=30.72 MHz if both channels are enabled - rx_lo = int(args.fc) #4e9 + fs = int(args.fs) # must be <=30.72 MHz if both channels are enabled + rx_lo = int(args.fc) # 4e9 tx_lo = rx_lo # setup receive - rx_mode = args.rx_mode + rx_mode = args.rx_mode rx_gain = args.rx_gain - rx_n=args.rx_n + rx_n = args.rx_n - sdr_receiver = adi.ad9361(uri='ip:%s' % args.receiver_ip) + sdr_receiver = adi.ad9361(uri="ip:%s" % args.receiver_ip) sdr_receiver.rx_enabled_channels = [0, 1] sdr_receiver.sample_rate = fs - assert(sdr_receiver.sample_rate==fs) - sdr_receiver.rx_rf_bandwidth = int(3*args.fi) #fc0*5) #TODO! + assert sdr_receiver.sample_rate == fs + sdr_receiver.rx_rf_bandwidth = int(3 * args.fi) # fc0*5) #TODO! sdr_receiver.rx_lo = int(rx_lo) sdr_receiver.gain_control_mode = rx_mode sdr_receiver.rx_hardwaregain_chan0 = int(rx_gain) sdr_receiver.rx_hardwaregain_chan1 = int(rx_gain) sdr_receiver.rx_buffer_size = int(rx_n) - sdr_receiver._rxadc.set_kernel_buffers_count(2) # set buffers to 1 (instead of the default 4) to avoid stale data on Pluto + sdr_receiver._rxadc.set_kernel_buffers_count( + 2 + ) # set buffers to 1 (instead of the default 4) to avoid stale data on Pluto - #drop the first bunch of frames + # drop the first bunch of frames for _ in range(20): sdr_receiver.rx() - freq = np.fft.fftfreq(rx_n,d=1.0/fs) + freq = np.fft.fftfreq(rx_n, d=1.0 / fs) - emitter_fs=16e6 + emitter_fs = 16e6 - retries=0 - while retries<10: - #try to setup TX and lets see if it works - sdr_emitter = adi.ad9361(uri='ip:%s' % args.emitter_ip) + retries = 0 + while retries < 10: + # try to setup TX and lets see if it works + sdr_emitter = adi.ad9361(uri="ip:%s" % args.emitter_ip) - #setup TX - sdr_emitter.sample_rate=emitter_fs - assert(sdr_emitter.sample_rate==emitter_fs) - sdr_emitter.tx_rf_bandwidth = int(3*args.fi) - assert(sdr_emitter.tx_rf_bandwidth==int(3*args.fi)) + # setup TX + sdr_emitter.sample_rate = emitter_fs + assert sdr_emitter.sample_rate == emitter_fs + sdr_emitter.tx_rf_bandwidth = int(3 * args.fi) + assert sdr_emitter.tx_rf_bandwidth == int(3 * args.fi) sdr_emitter.tx_lo = int(tx_lo) - assert(sdr_emitter.tx_lo==tx_lo) + assert sdr_emitter.tx_lo == tx_lo sdr_emitter.tx_enabled_channels = [0] - sdr_emitter.tx_hardwaregain_chan0 = int(args.tx_gain) #tx_gain) #tx_gain) - assert(sdr_emitter.tx_hardwaregain_chan0==int(args.tx_gain)) - sdr_emitter.tx_hardwaregain_chan1 = int(-80) # use Tx2 for calibration + sdr_emitter.tx_hardwaregain_chan0 = int(args.tx_gain) # tx_gain) #tx_gain) + assert sdr_emitter.tx_hardwaregain_chan0 == int(args.tx_gain) + sdr_emitter.tx_hardwaregain_chan1 = int(-80) # use Tx2 for calibration # - tx_n=int(fs/gcd(fs,fc0)) - while tx_n<1024*16: - tx_n*=2 - #sdr_emitter.tx_buffer_size = tx_n - - #since its a cyclic buffer its important to end on a full phase - t = np.arange(0, tx_n)/fs # time at each point assuming we are sending samples at (1/fs)s - iq0 = np.exp(1j*2*np.pi*fc0*t)*(2**14) - #try to reset the tx - #sdr_emitter.tx_destroy_buffer() + tx_n = int(fs / gcd(fs, fc0)) + while tx_n < 1024 * 16: + tx_n *= 2 + # sdr_emitter.tx_buffer_size = tx_n + + # since its a cyclic buffer its important to end on a full phase + t = ( + np.arange(0, tx_n) / fs + ) # time at each point assuming we are sending samples at (1/fs)s + iq0 = np.exp(1j * 2 * np.pi * fc0 * t) * (2**14) + # try to reset the tx + # sdr_emitter.tx_destroy_buffer() sdr_emitter.tx_destroy_buffer() - sdr_emitter.tx_cyclic_buffer = True # this keeps repeating! - assert(sdr_emitter.tx_cyclic_buffer==True) + sdr_emitter.tx_cyclic_buffer = True # this keeps repeating! + assert sdr_emitter.tx_cyclic_buffer == True sdr_emitter.tx(iq0) # Send Tx data. - - #give RX a chance to calm down + + # give RX a chance to calm down for _ in range(50): sdr_receiver.rx() - #test to see what frequency we are seeing - signal_matrix=np.vstack(sdr_receiver.rx()) + # test to see what frequency we are seeing + signal_matrix = np.vstack(sdr_receiver.rx()) sp = np.fft.fft(signal_matrix[0]) - max_freq=freq[np.abs(np.argmax(sp.real))] - if np.abs(max_freq-args.fi)<(args.fs/rx_n+1): + max_freq = freq[np.abs(np.argmax(sp.real))] + if np.abs(max_freq - args.fi) < (args.fs / rx_n + 1): print("Emitter online after %d retries" % retries) - return sdr_receiver,sdr_emitter - retries+=1 - sdr_emitter=None + return sdr_receiver, sdr_emitter + retries += 1 + sdr_emitter = None time.sleep(1) - return None,None + return None, None + -def circular_mean(angles,trim=50.0): - cm=np.arctan2(np.sin(angles).sum(),np.cos(angles).sum())%(2*np.pi) - dists=np.vstack([2*np.pi-np.abs(cm-angles),np.abs(cm-angles)]).min(axis=0) - _angles=angles[dists50: - #print(time.time(),sds.max(),thetas[sds.argmax()]) - counts[sds.argmax()%(intervals-1)]+=1 - - axs[1].cla() - axs[1].set_xlabel("Time") - axs[1].set_ylabel("Value") - axs[1].scatter(np.arange(signal_matrix.shape[1]),signal_matrix[0].real,s=2,label="I") - axs[1].scatter(np.arange(signal_matrix.shape[1]),signal_matrix[0].imag,s=2,label="Q") - axs[1].set_title("Receiver 1") - axs[1].legend(loc=3) - - axs[2].cla() - axs[2].set_xlabel("Time") - axs[2].set_ylabel("Value") - #axs[2].scatter(np.arange(signal_matrix.shape[1]),signal_matrix[1].real,s=2,label="I") - #axs[2].scatter(np.arange(signal_matrix.shape[1]),signal_matrix[1].imag,s=2,label="Q") - axs[2].scatter(signal_matrix[0].real,signal_matrix[0].imag,s=2,alpha=0.5,label="RX0") - axs[2].scatter(signal_matrix[1].real,signal_matrix[1].imag,s=2,alpha=0.5,label="RX1") - axs[2].legend(loc=3) - - axs[0].cla() - axs[0].stairs(counts, 360*thetas/(2*np.pi)) - axs[0].set_title("Histogram") - axs[0].set_xlabel("Theta") - axs[0].set_ylabel("Count") - - axs[3].cla() - axs[3].scatter(360*thetas/(2*np.pi),sds,s=0.5) - axs[3].set_title("Beamformer") - axs[3].set_xlabel("Theta") - axs[3].set_ylabel("Signal") - - sp = np.fft.fft(signal_matrix[0]) - axs[4].cla() - axs[4].set_title("FFT") - axs[4].scatter(freq, sp.real,s=1) #, freq, sp.imag) - max_freq=freq[np.abs(np.argmax(sp.real))] - axs[4].axvline( - x=max_freq, - label="max %0.2e" % max_freq, - color='red' - ) - axs[4].legend(loc=3) - #print("MAXFREQ",freq[np.abs(np.argmax(sp.real))]) - plt.draw() - plt.pause(0.01) + signal_matrix = np.vstack(sdr.rx()) + signal_matrix[0] *= np.exp(-1j * (2 * np.pi / 360) * args.cal0) + thetas, sds, steering = beamformer( + detector.all_receiver_pos(), signal_matrix, args.fc, spacing=intervals + ) + if sds.max() > 50: + # print(time.time(),sds.max(),thetas[sds.argmax()]) + counts[sds.argmax() % (intervals - 1)] += 1 + + axs[1].cla() + axs[1].set_xlabel("Time") + axs[1].set_ylabel("Value") + axs[1].scatter( + np.arange(signal_matrix.shape[1]), signal_matrix[0].real, s=2, label="I" + ) + axs[1].scatter( + np.arange(signal_matrix.shape[1]), signal_matrix[0].imag, s=2, label="Q" + ) + axs[1].set_title("Receiver 1") + axs[1].legend(loc=3) + + axs[2].cla() + axs[2].set_xlabel("Time") + axs[2].set_ylabel("Value") + # axs[2].scatter(np.arange(signal_matrix.shape[1]),signal_matrix[1].real,s=2,label="I") + # axs[2].scatter(np.arange(signal_matrix.shape[1]),signal_matrix[1].imag,s=2,label="Q") + axs[2].scatter( + signal_matrix[0].real, signal_matrix[0].imag, s=2, alpha=0.5, label="RX0" + ) + axs[2].scatter( + signal_matrix[1].real, signal_matrix[1].imag, s=2, alpha=0.5, label="RX1" + ) + axs[2].legend(loc=3) + + axs[0].cla() + axs[0].stairs(counts, 360 * thetas / (2 * np.pi)) + axs[0].set_title("Histogram") + axs[0].set_xlabel("Theta") + axs[0].set_ylabel("Count") + + axs[3].cla() + axs[3].scatter(360 * thetas / (2 * np.pi), sds, s=0.5) + axs[3].set_title("Beamformer") + axs[3].set_xlabel("Theta") + axs[3].set_ylabel("Signal") + + sp = np.fft.fft(signal_matrix[0]) + axs[4].cla() + axs[4].set_title("FFT") + axs[4].scatter(freq, sp.real, s=1) # , freq, sp.imag) + max_freq = freq[np.abs(np.argmax(sp.real))] + axs[4].axvline(x=max_freq, label="max %0.2e" % max_freq, color="red") + axs[4].legend(loc=3) + # print("MAXFREQ",freq[np.abs(np.argmax(sp.real))]) + plt.draw() + plt.pause(0.01) diff --git a/software/sdrpluto/sdr.py b/software/sdrpluto/sdr.py index e4bee7ac..fc6b1fb9 100644 --- a/software/sdrpluto/sdr.py +++ b/software/sdrpluto/sdr.py @@ -1,173 +1,230 @@ import matplotlib.pyplot as plt import numpy as np -''' +""" Given some guess of the source of direction we can shift the carrier frequency phase of received samples at the N different receivers. If the guess of the source direction is correct, the signal from the N different receivers should interfer constructively. -''' +""" + +c = 3e8 # speed of light -c=3e8 # speed of light class Source(object): - def __init__(self,pos): - self.pos=np.array(pos) + def __init__(self, pos): + self.pos = np.array(pos) + + def signal(self, sampling_times): + return ( + np.cos(2 * np.pi * sampling_times) + np.sin(2 * np.pi * sampling_times) * 1j + ) - def signal(self,sampling_times): - return np.cos(2*np.pi*sampling_times)+np.sin(2*np.pi*sampling_times)*1j + def demod_signal(self, signal, demod_times): + return signal - def demod_signal(self,signal,demod_times): - return signal class SinSource(Source): - def __init__(self,pos,frequency,phase): - super().__init__(pos) - self.frequency=frequency - self.phase=phase + def __init__(self, pos, frequency, phase): + super().__init__(pos) + self.frequency = frequency + self.phase = phase + + def signal(self, sampling_times): + return ( + np.cos(2 * np.pi * sampling_times * self.frequency + self.phase) + + np.sin(2 * np.pi * sampling_times * self.frequency + self.phase) * 1j + ) - def signal(self,sampling_times): - return np.cos(2*np.pi*sampling_times*self.frequency+self.phase)+np.sin(2*np.pi*sampling_times*self.frequency+self.phase)*1j class MixedSource(Source): - def __init__(self,source_a,source_b): - super().__init__(pos) - self.source_a=source_a - self.source_b=source_b + def __init__(self, source_a, source_b): + super().__init__(pos) + self.source_a = source_a + self.source_b = source_b + + def signal(self, sampling_times): + return self.source_a(sampling_times) * self.source_b(sampling_times) - def signal(self,sampling_times): - return self.source_a(sampling_times)*self.source_b(sampling_times) class QAMSource(Source): - def __init__(self,pos,carrier_frequency,signal_frequency,sigma=0): - super().__init__(pos) - self.lo_in_phase=SinSource(pos,carrier_frequency,-np.pi/2) # cos - self.lo_out_of_phase=SinSource(pos,carrier_frequency,0) # cos - self.signal_source=SinSource(pos,signal_frequency,0) - self.sigma=sigma - - def signal(self,sampling_times): - signal=self.signal_source.signal(sampling_times) - return ((self.lo_in_phase.signal(sampling_times)*signal.real+\ - self.lo_out_of_phase.signal(sampling_times)*signal.imag)/2) +\ - np.random.randn(sampling_times.shape[0])*self.sigma - - def demod_signal(self,signal,demod_times): - return (self.lo_in_phase(demod_times)+\ - self.lo_out_of_phase(demod_times)*1j)*signal - + def __init__(self, pos, carrier_frequency, signal_frequency, sigma=0): + super().__init__(pos) + self.lo_in_phase = SinSource(pos, carrier_frequency, -np.pi / 2) # cos + self.lo_out_of_phase = SinSource(pos, carrier_frequency, 0) # cos + self.signal_source = SinSource(pos, signal_frequency, 0) + self.sigma = sigma + + def signal(self, sampling_times): + signal = self.signal_source.signal(sampling_times) + return ( + ( + self.lo_in_phase.signal(sampling_times) * signal.real + + self.lo_out_of_phase.signal(sampling_times) * signal.imag + ) + / 2 + ) + np.random.randn(sampling_times.shape[0]) * self.sigma + + def demod_signal(self, signal, demod_times): + return ( + self.lo_in_phase(demod_times) + self.lo_out_of_phase(demod_times) * 1j + ) * signal + class NoiseWrapper(Source): - def __init__(self,internal_source,sigma=1): - super().__init__(internal_source.pos) - self.internal_source=internal_source - self.sigma=sigma + def __init__(self, internal_source, sigma=1): + super().__init__(internal_source.pos) + self.internal_source = internal_source + self.sigma = sigma + + def signal(self, sampling_times): + return self.internal_source.signal(sampling_times) + ( + np.random.randn(sampling_times.shape[0], 2) * self.sigma + ).view(np.cdouble).reshape(-1) - def signal(self,sampling_times): - return self.internal_source.signal(sampling_times) + (np.random.randn(sampling_times.shape[0], 2)*self.sigma).view(np.cdouble).reshape(-1) class Receiver: - def __init__(self,pos): - self.pos=np.array(pos) + def __init__(self, pos): + self.pos = np.array(pos) + class Detector(object): - def __init__(self,sampling_frequency): - self.sources=[] - self.receivers=[] - self.sampling_frequency=sampling_frequency - - def add_source(self,source): - self.sources.append(source) - - def rm_sources(self): - self.sources=[] - - def add_receiver(self,receiver): - self.receivers.append(receiver) - - def get_signal_matrix(self,start_time,duration,rx_lo=0): - n_samples=int(duration*self.sampling_frequency) - base_times=start_time+np.linspace(0,n_samples-1,n_samples)/self.sampling_frequency - sample_matrix=np.zeros((len(self.receivers),n_samples),dtype=np.cdouble) # receivers x samples - for receiver_index,receiver in enumerate(self.receivers): - for _source in self.sources: - time_delay=np.linalg.norm(receiver.pos-_source.pos)/c - sample_matrix[receiver_index,:]+=_source.demod_signal(_source.signal(base_times-time_delay), - base_times) - if rx_lo>0: - sample_matrix[receiver_index,:] - return sample_matrix - - -def beamformer(detector,signal_matrix,carrier_frequency,calibration=None,spacing=64+1): - if calibration is None: - calibration=np.ones(len(detector.receivers)).astype(np.cdouble) - thetas=np.linspace(-np.pi,np.pi,spacing) - steer_dot_signal=np.zeros(thetas.shape[0]) - carrier_wavelength=c/carrier_frequency - steering_vectors=np.zeros((len(thetas),len(detector.receivers))).astype(np.cdouble) - for theta_index,theta in enumerate(thetas): - source_vector=np.array([np.cos(theta),np.sin(theta)]) - projections=[] - for receiver_index,receiver in enumerate(detector.receivers): - projection_of_receiver_onto_source_direction=np.dot(source_vector,receiver.pos) - projections.append(projection_of_receiver_onto_source_direction/carrier_wavelength) - arg=2*np.pi*projection_of_receiver_onto_source_direction/carrier_wavelength - steering_vectors[theta_index][receiver_index]=np.exp(-1j*arg) - steer_dot_signal[theta_index]=np.absolute(np.matmul(steering_vectors[theta_index]*calibration,signal_matrix)).mean() - return thetas,steer_dot_signal,steering_vectors - -def plot_space(ax,d,wavelength=1): - #fig,ax=plt.subplots(1,1,figsize=(4,4)) - receiver_pos=np.vstack( - [ receiver.pos/wavelength for receiver in d.receivers ] + def __init__(self, sampling_frequency): + self.sources = [] + self.receivers = [] + self.sampling_frequency = sampling_frequency + + def add_source(self, source): + self.sources.append(source) + + def rm_sources(self): + self.sources = [] + + def add_receiver(self, receiver): + self.receivers.append(receiver) + + def get_signal_matrix(self, start_time, duration, rx_lo=0): + n_samples = int(duration * self.sampling_frequency) + base_times = ( + start_time + + np.linspace(0, n_samples - 1, n_samples) / self.sampling_frequency + ) + sample_matrix = np.zeros( + (len(self.receivers), n_samples), dtype=np.cdouble + ) # receivers x samples + for receiver_index, receiver in enumerate(self.receivers): + for _source in self.sources: + time_delay = np.linalg.norm(receiver.pos - _source.pos) / c + sample_matrix[receiver_index, :] += _source.demod_signal( + _source.signal(base_times - time_delay), base_times ) - _max=receiver_pos.max() - _min=receiver_pos.min() - buffer=(_max-_min)*0.1 - _max+=buffer - _min-=buffer + if rx_lo > 0: + sample_matrix[receiver_index, :] + return sample_matrix + - center_mass=receiver_pos.mean(axis=0) +def beamformer( + detector, signal_matrix, carrier_frequency, calibration=None, spacing=64 + 1 +): + if calibration is None: + calibration = np.ones(len(detector.receivers)).astype(np.cdouble) + thetas = np.linspace(-np.pi, np.pi, spacing) + steer_dot_signal = np.zeros(thetas.shape[0]) + carrier_wavelength = c / carrier_frequency + steering_vectors = np.zeros((len(thetas), len(detector.receivers))).astype( + np.cdouble + ) + for theta_index, theta in enumerate(thetas): + source_vector = np.array([np.cos(theta), np.sin(theta)]) + projections = [] + for receiver_index, receiver in enumerate(detector.receivers): + projection_of_receiver_onto_source_direction = np.dot( + source_vector, receiver.pos + ) + projections.append( + projection_of_receiver_onto_source_direction / carrier_wavelength + ) + arg = ( + 2 + * np.pi + * projection_of_receiver_onto_source_direction + / carrier_wavelength + ) + steering_vectors[theta_index][receiver_index] = np.exp(-1j * arg) + steer_dot_signal[theta_index] = np.absolute( + np.matmul(steering_vectors[theta_index] * calibration, signal_matrix) + ).mean() + return thetas, steer_dot_signal, steering_vectors + + +def plot_space(ax, d, wavelength=1): + # fig,ax=plt.subplots(1,1,figsize=(4,4)) + receiver_pos = np.vstack([receiver.pos / wavelength for receiver in d.receivers]) + _max = receiver_pos.max() + _min = receiver_pos.min() + buffer = (_max - _min) * 0.1 + _max += buffer + _min -= buffer + + center_mass = receiver_pos.mean(axis=0) + + source_vectors = [ + (source.pos / wavelength - center_mass) + / np.linalg.norm(source.pos / wavelength - center_mass) + for source in d.sources + ] + + ax.set_xlim([_min, _max]) + ax.set_ylim([_min, _max]) + + ax.scatter(receiver_pos[:, 0], receiver_pos[:, 1], label="Receivers") + thetas = [] + for source_vector in source_vectors: + ax.quiver( + center_mass[0], + center_mass[1], + -source_vector[0], + -source_vector[1], + scale=5, + alpha=0.5, + color="red", + label="Source", + ) + thetas.append( + 360 * np.arctan2(source_vector[1], source_vector[0]) / (2 * np.pi) + ) + + ax.legend() + ax.set_xlabel("x (wavelengths)") + ax.set_ylabel("y (wavelengths)") + ax.set_title("Space diagram (%s)" % ",".join(map(lambda x: "%0.2f" % x, thetas))) - source_vectors=[ (source.pos/wavelength-center_mass)/np.linalg.norm(source.pos/wavelength-center_mass) for source in d.sources ] - ax.set_xlim([_min,_max]) - ax.set_ylim([_min,_max]) +class ULADetector(Detector): + def __init__(self, sampling_frequency, n_elements, spacing): + super().__init__(sampling_frequency) + for idx in np.arange(n_elements): + self.add_receiver(Receiver([spacing * (idx - (n_elements - 1) / 2), 0])) - ax.scatter(receiver_pos[:,0],receiver_pos[:,1],label="Receivers") - thetas=[] - for source_vector in source_vectors: - ax.quiver(center_mass[0], center_mass[1], - -source_vector[0], -source_vector[1], scale=5, alpha=0.5,color='red',label="Source") - thetas.append( 360*np.arctan2(source_vector[1],source_vector[0])/(2*np.pi) ) - ax.legend() - ax.set_xlabel("x (wavelengths)") - ax.set_ylabel("y (wavelengths)") - ax.set_title("Space diagram (%s)" % ",".join(map(lambda x : "%0.2f" % x ,thetas))) +ula_d = ULADetector(300, 10, 1) +fig, ax = plt.subplots(1, 1) +plot_space(ax, ula_d) -class ULADetector(Detector): - def __init__(self,sampling_frequency,n_elements,spacing): - super().__init__(sampling_frequency) - for idx in np.arange(n_elements): - self.add_receiver(Receiver([ - spacing*(idx-(n_elements-1)/2), - 0])) -ula_d=ULADetector(300,10,1) -fig,ax=plt.subplots(1,1) -plot_space(ax,ula_d) - class UCADetector(Detector): - def __init__(self,sampling_frequency,n_elements,radius): - super().__init__(sampling_frequency) - for theta in np.linspace(0,2*np.pi,n_elements+1)[:-1]+np.pi/2: # orientate along y axis - self.add_receiver(Receiver([ - radius*np.cos(theta), - radius*np.sin(theta)])) -uca_d=UCADetector(300,10,1) -fig,ax=plt.subplots(1,1) -plot_space(ax,uca_d) + def __init__(self, sampling_frequency, n_elements, radius): + super().__init__(sampling_frequency) + for theta in ( + np.linspace(0, 2 * np.pi, n_elements + 1)[:-1] + np.pi / 2 + ): # orientate along y axis + self.add_receiver( + Receiver([radius * np.cos(theta), radius * np.sin(theta)]) + ) + + +uca_d = UCADetector(300, 10, 1) +fig, ax = plt.subplots(1, 1) +plot_space(ax, uca_d) diff --git a/software/sdrpluto/test.py b/software/sdrpluto/test.py index a4198323..df8256ed 100644 --- a/software/sdrpluto/test.py +++ b/software/sdrpluto/test.py @@ -1,7 +1,8 @@ import iio -ctx = iio.Context('ip:pluto.local') + +ctx = iio.Context("ip:pluto.local") for dev in ctx.devices: for chan in dev.channels: for attr in chan.attrs: - print(dev.name,chan.name,attr) - #print(f'{dev.name}: {chan.name}: {attr.name}') + print(dev.name, chan.name, attr) + # print(f'{dev.name}: {chan.name}: {attr.name}') diff --git a/software/sdrpluto/test_emitter_recv/03_only_emit.py b/software/sdrpluto/test_emitter_recv/03_only_emit.py index 943cda65..90710c2d 100644 --- a/software/sdrpluto/test_emitter_recv/03_only_emit.py +++ b/software/sdrpluto/test_emitter_recv/03_only_emit.py @@ -1,60 +1,69 @@ -#Starter code from jon Kraft +# Starter code from jon Kraft import adi import numpy as np -from math import lcm,gcd +from math import lcm, gcd import argparse parser = argparse.ArgumentParser() -parser.add_argument("--ip", type=str, help="target Pluto IP address",required=True) -parser.add_argument("--fi", type=int, help="Intermediate frequency",required=False,default=1e5) -parser.add_argument("--fc", type=int, help="Intermediate frequency",required=False,default=2.5e9) -parser.add_argument("--fs", type=int, help="Intermediate frequency",required=False,default=28e6) +parser.add_argument("--ip", type=str, help="target Pluto IP address", required=True) +parser.add_argument( + "--fi", type=int, help="Intermediate frequency", required=False, default=1e5 +) +parser.add_argument( + "--fc", type=int, help="Intermediate frequency", required=False, default=2.5e9 +) +parser.add_argument( + "--fs", type=int, help="Intermediate frequency", required=False, default=28e6 +) args = parser.parse_args() -c=3e8 +c = 3e8 fc0 = int(args.fi) -fs = int(args.fs) # must be <=30.72 MHz if both channels are enabled -rx_lo = int(args.fc) #4e9 +fs = int(args.fs) # must be <=30.72 MHz if both channels are enabled +rx_lo = int(args.fc) # 4e9 tx_lo = rx_lo rx_mode = "manual" # can be "manual" or "slow_attack" tx_gain = -3 -wavelength = c/rx_lo # wavelength of the RF carrier +wavelength = c / rx_lo # wavelength of the RF carrier -rx_n=int(2**10) +rx_n = int(2**10) -sdr = adi.ad9361(uri='ip:%s' % args.ip) +sdr = adi.ad9361(uri="ip:%s" % args.ip) -#setup TX -sdr.sample_rate=fs +# setup TX +sdr.sample_rate = fs sdr.tx_rf_bandwidth = int(fs) sdr.tx_lo = int(tx_lo) -sdr.tx_cyclic_buffer = True # this keeps repeating! -sdr.tx_hardwaregain_chan0 = int(-5) #tx_gain) #tx_gain) -sdr.tx_hardwaregain_chan1 = int(-80) # use Tx2 for calibration +sdr.tx_cyclic_buffer = True # this keeps repeating! +sdr.tx_hardwaregain_chan0 = int(-5) # tx_gain) #tx_gain) +sdr.tx_hardwaregain_chan1 = int(-80) # use Tx2 for calibration # -#tx_n=int(min(lcm(fc0,fs),rx_n*8)) #1024*1024*1024) # tx for longer than rx -#tx_n=int(lcm(fc0,fs)) -#tx_n=fc0 #*8 -tx_n=int(fs/gcd(fs,fc0)) -#tx_n=int(lcm(fc0,fs)) -while tx_n<1024*16: - tx_n*=2 -print("TX_N",tx_n,fs,fc0) +# tx_n=int(min(lcm(fc0,fs),rx_n*8)) #1024*1024*1024) # tx for longer than rx +# tx_n=int(lcm(fc0,fs)) +# tx_n=fc0 #*8 +tx_n = int(fs / gcd(fs, fc0)) +# tx_n=int(lcm(fc0,fs)) +while tx_n < 1024 * 16: + tx_n *= 2 +print("TX_N", tx_n, fs, fc0) sdr.tx_buffer_size = tx_n -print(sdr.tx_lo,tx_n,lcm(fc0,fs)) +print(sdr.tx_lo, tx_n, lcm(fc0, fs)) -#since its a cyclic buffer its important to end on a full phase -t = np.arange(0, tx_n)/fs # time at each point assuming we are sending samples at (1/fs)s -iq0 = np.exp(1j*2*np.pi*fc0*t)*(2**14) +# since its a cyclic buffer its important to end on a full phase +t = ( + np.arange(0, tx_n) / fs +) # time at each point assuming we are sending samples at (1/fs)s +iq0 = np.exp(1j * 2 * np.pi * fc0 * t) * (2**14) print(iq0) sdr.tx_enabled_channels = [0] sdr.tx(iq0) # Send Tx data. import time + time.sleep(1) time.sleep(1) print("SLEEP") diff --git a/software/sdrpluto/test_emitter_recv/04_plot_signal.py b/software/sdrpluto/test_emitter_recv/04_plot_signal.py index afe4286f..a566fd9a 100644 --- a/software/sdrpluto/test_emitter_recv/04_plot_signal.py +++ b/software/sdrpluto/test_emitter_recv/04_plot_signal.py @@ -1,75 +1,76 @@ -#Starter code from jon Kraft +# Starter code from jon Kraft import time import adi import matplotlib.pyplot as plt import numpy as np from math import lcm import argparse -#from sdr import * + +# from sdr import * parser = argparse.ArgumentParser() -parser.add_argument("--ip", type=str, help="target Pluto IP address",required=True) -parser.add_argument("--fi", type=int, help="Intermediate frequency",required=False,default=1e5) +parser.add_argument("--ip", type=str, help="target Pluto IP address", required=True) +parser.add_argument( + "--fi", type=int, help="Intermediate frequency", required=False, default=1e5 +) args = parser.parse_args() -c=3e8 +c = 3e8 fc0 = int(args.fi) -fs = int(4e6) # must be <=30.72 MHz if both channels are enabled -rx_lo = int(2.5e9) #4e9 +fs = int(4e6) # must be <=30.72 MHz if both channels are enabled +rx_lo = int(2.5e9) # 4e9 tx_lo = rx_lo rx_mode = "manual" # can be "manual" or "slow_attack" -rx_gain = -10 +rx_gain = -10 tx_gain = -30 -wavelength = c/rx_lo # wavelength of the RF carrier +wavelength = c / rx_lo # wavelength of the RF carrier -rx_n=int(2**14) +rx_n = int(2**14) -sdr = adi.ad9361(uri='ip:%s' % args.ip) +sdr = adi.ad9361(uri="ip:%s" % args.ip) -#setup RX +# setup RX sdr.rx_enabled_channels = [0, 1] sdr.sample_rate = fs -assert(sdr.sample_rate==fs) -sdr.rx_rf_bandwidth = int(fs) #fc0*5) +assert sdr.sample_rate == fs +sdr.rx_rf_bandwidth = int(fs) # fc0*5) sdr.rx_lo = int(rx_lo) sdr.gain_control_mode = rx_mode sdr.rx_hardwaregain_chan0 = int(rx_gain) sdr.rx_hardwaregain_chan1 = int(rx_gain) sdr.rx_buffer_size = int(rx_n) -sdr._rxadc.set_kernel_buffers_count(1) # set buffers to 1 (instead of the default 4) to avoid stale data on Pluto +sdr._rxadc.set_kernel_buffers_count( + 1 +) # set buffers to 1 (instead of the default 4) to avoid stale data on Pluto -#setup TX +# setup TX sdr.tx_enabled_channels = [] -#sdr.tx_rf_bandwidth = int(fc0*3) -#sdr.tx_lo = int(tx_lo) -#sdr.tx_cyclic_buffer = True # this keeps repeating! -#sdr.tx_hardwaregain_chan0 = int(-80) #tx_gain) #tx_gain) -#sdr.tx_hardwaregain_chan1 = int(-80) # use Tx2 for calibration -#fig,axs=plt.subplots(1,1,figsize=(4,4)) +# sdr.tx_rf_bandwidth = int(fc0*3) +# sdr.tx_lo = int(tx_lo) +# sdr.tx_cyclic_buffer = True # this keeps repeating! +# sdr.tx_hardwaregain_chan0 = int(-80) #tx_gain) #tx_gain) +# sdr.tx_hardwaregain_chan1 = int(-80) # use Tx2 for calibration +# fig,axs=plt.subplots(1,1,figsize=(4,4)) -fig,axs=plt.subplots(2,2,figsize=(8,6)) +fig, axs = plt.subplots(2, 2, figsize=(8, 6)) -t=np.arange(rx_n) +t = np.arange(rx_n) while True: - signal_matrix=np.vstack(sdr.rx()) + signal_matrix = np.vstack(sdr.rx()) - freq = np.fft.fftfreq(t.shape[-1],d=1.0/fs) - assert(t.shape[-1]==rx_n) - for idx in [0,1]: - axs[idx][0].clear() - axs[idx][1].clear() - axs[idx][0].scatter(t,signal_matrix[idx].real,s=1) - sp = np.fft.fft(signal_matrix[idx]) - axs[idx][1].scatter(freq, sp.real,s=1) #, freq, sp.imag) - max_freq=freq[np.abs(np.argmax(sp.real))] - axs[idx][1].axvline( - x=max_freq, - label="max %0.2e" % max_freq, - color='red' - ) - axs[idx][1].legend() - print("MAXFREQ",freq[np.abs(np.argmax(sp.real))]) - fig.canvas.draw() - plt.pause(0.00001) + freq = np.fft.fftfreq(t.shape[-1], d=1.0 / fs) + assert t.shape[-1] == rx_n + for idx in [0, 1]: + axs[idx][0].clear() + axs[idx][1].clear() + axs[idx][0].scatter(t, signal_matrix[idx].real, s=1) + sp = np.fft.fft(signal_matrix[idx]) + axs[idx][1].scatter(freq, sp.real, s=1) # , freq, sp.imag) + max_freq = freq[np.abs(np.argmax(sp.real))] + axs[idx][1].axvline(x=max_freq, label="max %0.2e" % max_freq, color="red") + axs[idx][1].legend() + print("MAXFREQ", freq[np.abs(np.argmax(sp.real))]) + fig.canvas.draw() + plt.pause(0.00001) diff --git a/software/sdrpluto/test_single_phase/02_wifi_direction.py b/software/sdrpluto/test_single_phase/02_wifi_direction.py index bac25517..8ca73b4b 100644 --- a/software/sdrpluto/test_single_phase/02_wifi_direction.py +++ b/software/sdrpluto/test_single_phase/02_wifi_direction.py @@ -1,111 +1,125 @@ -#Starter code from jon Kraft +# Starter code from jon Kraft import adi import matplotlib.pyplot as plt import numpy as np -from math import lcm,gcd +from math import lcm, gcd import argparse from utils.rf import ULADetector, beamformer, beamformer_old, dbfs parser = argparse.ArgumentParser() -parser.add_argument("--ip", type=str, help="target Pluto IP address",required=True) -parser.add_argument("--fi", type=int, help="Intermediate frequency",required=False,default=1e5) -parser.add_argument("--fc", type=int, help="Intermediate frequency",required=False,default=2.5e9) -parser.add_argument("--fs", type=int, help="Intermediate frequency",required=False,default=16e6) -parser.add_argument("--cal0", type=int, help="Rx0 calibration phase offset in degrees",required=False,default=180) -parser.add_argument("--d", type=int, help="Distance apart",required=False,default=0.062) +parser.add_argument("--ip", type=str, help="target Pluto IP address", required=True) +parser.add_argument( + "--fi", type=int, help="Intermediate frequency", required=False, default=1e5 +) +parser.add_argument( + "--fc", type=int, help="Intermediate frequency", required=False, default=2.5e9 +) +parser.add_argument( + "--fs", type=int, help="Intermediate frequency", required=False, default=16e6 +) +parser.add_argument( + "--cal0", + type=int, + help="Rx0 calibration phase offset in degrees", + required=False, + default=180, +) +parser.add_argument( + "--d", type=int, help="Distance apart", required=False, default=0.062 +) args = parser.parse_args() - -c=3e8 +c = 3e8 fc0 = int(args.fi) -fs = int(args.fs) # must be <=30.72 MHz if both channels are enabled -rx_lo = int(args.fc) #4e9 +fs = int(args.fs) # must be <=30.72 MHz if both channels are enabled +rx_lo = int(args.fc) # 4e9 tx_lo = rx_lo rx_mode = "slow_attack" # can be "manual" or "slow_attack" -rx_gain = 40 +rx_gain = 40 tx_gain = -30 -rx_n=int(2**15) +rx_n = int(2**15) -detector=ULADetector(fs,2,args.d) +detector = ULADetector(fs, 2, args.d) -sdr = adi.ad9361(uri='ip:%s' % args.ip) +sdr = adi.ad9361(uri="ip:%s" % args.ip) -#setup RX +# setup RX sdr.rx_enabled_channels = [0, 1] sdr.sample_rate = fs -assert(sdr.sample_rate==fs) -sdr.rx_rf_bandwidth = int(fc0*3) +assert sdr.sample_rate == fs +sdr.rx_rf_bandwidth = int(fc0 * 3) sdr.rx_lo = int(rx_lo) sdr.gain_control_mode = rx_mode sdr.rx_hardwaregain_chan0 = int(rx_gain) sdr.rx_hardwaregain_chan1 = int(rx_gain) sdr.rx_buffer_size = int(rx_n) -sdr._rxadc.set_kernel_buffers_count(1) # set buffers to 1 (instead of the default 4) to avoid stale data on Pluto +sdr._rxadc.set_kernel_buffers_count( + 1 +) # set buffers to 1 (instead of the default 4) to avoid stale data on Pluto -#setup TX -#sdr.tx_enabled_channels = [] +# setup TX +# sdr.tx_enabled_channels = [] import time -fig,axs=plt.subplots(1,4,figsize=(16,4)) -intervals=2*64+1 -counts=np.zeros(64) +fig, axs = plt.subplots(1, 4, figsize=(16, 4)) + +intervals = 2 * 64 + 1 +counts = np.zeros(64) -freq = np.fft.fftfreq(rx_n,d=1.0/fs) +freq = np.fft.fftfreq(rx_n, d=1.0 / fs) while True: - signal_matrix=np.vstack(sdr.rx()) - signal_matrix[0]*=np.exp(-1j*(2*np.pi/360)*args.cal0) - thetas,sds,steering=beamformer( - detector.all_receiver_pos(), - signal_matrix, - args.fc, - spacing=intervals) - thetas=thetas[64:] - sds=sds[64:] - steering=steering[64:] - if sds.max()>50: - #print(time.time(),sds.max(),thetas[sds.argmax()]) - counts[sds.argmax()%64]+=1 - - axs[1].cla() - axs[1].set_xlabel("Time") - axs[1].set_ylabel("Value") - #axs[2].scatter(np.arange(signal_matrix.shape[1]),signal_matrix[1].real,s=2,label="I") - #axs[2].scatter(np.arange(signal_matrix.shape[1]),signal_matrix[1].imag,s=2,label="Q") - axs[1].scatter(signal_matrix[0].real,signal_matrix[0].imag,s=2,alpha=0.5,label="RX1") - axs[1].scatter(signal_matrix[1].real,signal_matrix[1].imag,s=2,alpha=0.5,label="RX2") - axs[1].set_title("Receiver signals") - axs[1].legend(loc=3) - - axs[0].cla() - axs[0].stairs(counts, 360*thetas/(2*np.pi)) - axs[0].set_title("Histogram") - axs[0].set_xlabel("Theta") - axs[0].set_ylabel("Count") - - axs[2].cla() - axs[2].scatter(360*thetas/(2*np.pi),sds,s=0.5) - axs[2].set_title("Beamformer") - axs[2].set_xlabel("Theta") - axs[2].set_ylabel("Signal") - - sp = np.fft.fft(signal_matrix[0]) - axs[3].cla() - axs[3].set_title("FFT") - axs[3].scatter(freq, sp.real,s=1) #, freq, sp.imag) - max_freq=freq[np.abs(np.argmax(sp.real))] - axs[3].axvline( - x=max_freq, - label="max %0.2e" % max_freq, - color='red' - ) - axs[3].legend(loc=3) - #print("MAXFREQ",freq[np.abs(np.argmax(sp.real))]) - plt.draw() - plt.pause(0.01) + signal_matrix = np.vstack(sdr.rx()) + signal_matrix[0] *= np.exp(-1j * (2 * np.pi / 360) * args.cal0) + thetas, sds, steering = beamformer( + detector.all_receiver_pos(), signal_matrix, args.fc, spacing=intervals + ) + thetas = thetas[64:] + sds = sds[64:] + steering = steering[64:] + if sds.max() > 50: + # print(time.time(),sds.max(),thetas[sds.argmax()]) + counts[sds.argmax() % 64] += 1 + + axs[1].cla() + axs[1].set_xlabel("Time") + axs[1].set_ylabel("Value") + # axs[2].scatter(np.arange(signal_matrix.shape[1]),signal_matrix[1].real,s=2,label="I") + # axs[2].scatter(np.arange(signal_matrix.shape[1]),signal_matrix[1].imag,s=2,label="Q") + axs[1].scatter( + signal_matrix[0].real, signal_matrix[0].imag, s=2, alpha=0.5, label="RX1" + ) + axs[1].scatter( + signal_matrix[1].real, signal_matrix[1].imag, s=2, alpha=0.5, label="RX2" + ) + axs[1].set_title("Receiver signals") + axs[1].legend(loc=3) + + axs[0].cla() + axs[0].stairs(counts, 360 * thetas / (2 * np.pi)) + axs[0].set_title("Histogram") + axs[0].set_xlabel("Theta") + axs[0].set_ylabel("Count") + + axs[2].cla() + axs[2].scatter(360 * thetas / (2 * np.pi), sds, s=0.5) + axs[2].set_title("Beamformer") + axs[2].set_xlabel("Theta") + axs[2].set_ylabel("Signal") + + sp = np.fft.fft(signal_matrix[0]) + axs[3].cla() + axs[3].set_title("FFT") + axs[3].scatter(freq, sp.real, s=1) # , freq, sp.imag) + max_freq = freq[np.abs(np.argmax(sp.real))] + axs[3].axvline(x=max_freq, label="max %0.2e" % max_freq, color="red") + axs[3].legend(loc=3) + # print("MAXFREQ",freq[np.abs(np.argmax(sp.real))]) + plt.draw() + plt.pause(0.01)