From d7d6c6d6ea62f9fcd88ef73177c5928d4186655c Mon Sep 17 00:00:00 2001 From: Misko Date: Tue, 19 Dec 2023 22:25:16 -0800 Subject: [PATCH] formatting --- hardware/3D_printing/antenna_mount.scad | 9 +- spf/baseline_algorithm.py | 15 +--- spf/grbl/grbl_interactive.py | 38 +++------ spf/grbl/mm_per_step.py | 4 +- spf/grbl_sdr_collect.py | 4 +- .../01_generate_data.py | 32 ++------ .../12_task2_model_training.py | 82 +++++-------------- .../13_learn_beamformer.py | 53 ++++-------- .../14_task3_model_training.py | 82 +++++-------------- .../90_real_session_plotter.py | 16 +--- .../91_line_plotter.py | 4 +- .../92_evaluate_session.py | 8 +- .../models/models.py | 61 ++++---------- spf/model_training_and_inference/test.py | 24 ++---- .../utils/image_utils.py | 22 ++--- .../utils/plot.py | 60 ++++++-------- .../utils/spf_dataset.py | 67 ++++----------- .../utils/spf_generate.py | 49 +++-------- spf/rf.py | 40 +++------ spf/sdrpluto/01_phase_sync.py | 8 +- spf/sdrpluto/02_wifi_direction.py | 8 +- spf/sdrpluto/gather.py | 40 +++------ spf/sdrpluto/phase_cal/01_phase_cal.py | 4 +- spf/sdrpluto/sdr.py | 28 ++----- .../test_emitter_recv/03_only_emit.py | 4 +- .../test_emitter_recv/04_plot_signal.py | 8 +- .../test_single_phase/02_wifi_direction.py | 4 +- 27 files changed, 214 insertions(+), 560 deletions(-) diff --git a/hardware/3D_printing/antenna_mount.scad b/hardware/3D_printing/antenna_mount.scad index 4d0b22c5..e28062c4 100644 --- a/hardware/3D_printing/antenna_mount.scad +++ b/hardware/3D_printing/antenna_mount.scad @@ -1,5 +1,5 @@ wave_length=125/2; -antenna_width=10; +antenna_width=10-0.2; antenna_stalk_width=10-0.25; antenna_base_length=23; @@ -9,7 +9,8 @@ height=antenna_width+8; ziptie_width=3.5; ziptie_depth=2.5; -screw_r=5.5/2; +dry_screw_r=5.5/2; //dry wall +screw_r=3.5/2; //m3 backing=2; edges=2; @@ -75,8 +76,8 @@ difference() { rotate([0,0,tx_theta]) antenna_cutout(); -k1=antenna_width/2+ziptie_depth+(screw_r+backing); -k2=wave_length/2+antenna_base_length-screw_r*2; +k1=antenna_width/2+ziptie_depth+(dry_screw_r+backing); +k2=wave_length/2+antenna_base_length-dry_screw_r*2; translate([k1,-k2,0]) screw_hold(); translate([-k1,k2,0]) screw_hold(); translate([k2,-k1,0]) screw_hold(); diff --git a/spf/baseline_algorithm.py b/spf/baseline_algorithm.py index 4e194f48..c47ce9a2 100644 --- a/spf/baseline_algorithm.py +++ b/spf/baseline_algorithm.py @@ -50,9 +50,7 @@ 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_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)) @@ -77,12 +75,9 @@ def get_top_n_points( ) if ( line_to_point_assignments[line_idx] == -1 - or distance_to_line - < line_to_point_distances[line_idx] + or distance_to_line < line_to_point_distances[line_idx] ): - line_to_point_assignments[line_idx] = ( - len(final_points) - 1 - ) + 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 @@ -179,9 +174,7 @@ def lines_to_points(lines, t): 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 - ): + 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: diff --git a/spf/grbl/grbl_interactive.py b/spf/grbl/grbl_interactive.py index 5087f3f4..7430825c 100644 --- a/spf/grbl/grbl_interactive.py +++ b/spf/grbl/grbl_interactive.py @@ -42,7 +42,7 @@ def a_to_b_in_stepsize(a, b, step_size): - if np.isclose(a,b).all(): + if np.isclose(a, b).all(): return [b] # move by step_size from where we are now to the target position points = [a] @@ -104,7 +104,7 @@ def to_steps(self, p): if (self.polygon is not None) and not self.polygon.contains_point( p, radius=0.01 ): # todo a bit hacky but works - print("OUT OF BOUNDS",p) + print("OUT OF BOUNDS", p) raise ValueError # motor_steps = ( distance_between_pivot and point ) - (distance between pivot and calibration point) a_motor_steps = np.linalg.norm(self.pA - p) - np.linalg.norm( @@ -123,13 +123,9 @@ def binary_search_edge(self, left, right, xy, direction, epsilon): try: steps = self.to_steps(p) # actual = self.from_steps(*steps) - return self.binary_search_edge( - midpoint, right, xy, direction, epsilon - ) + return self.binary_search_edge(midpoint, right, xy, direction, epsilon) except ValueError: - return self.binary_search_edge( - left, midpoint, xy, direction, epsilon - ) + return self.binary_search_edge(left, midpoint, xy, direction, epsilon) def get_boundary_vector_near_point(self, p): if self.polygon is None: @@ -167,13 +163,10 @@ def get_bounce_pos_and_new_direction(self, p, direction): # parallel component stays the same # negatate the perpendicular component - bvec = self.dynamics.get_boundary_vector_near_point( - last_point_before_bounce - ) + bvec = self.dynamics.get_boundary_vector_near_point(last_point_before_bounce) bvec_perp = np.array([bvec[1], -bvec[0]]) new_direction = ( - np.dot(direction, bvec) * bvec - - np.dot(direction, bvec_perp) * bvec_perp + np.dot(direction, bvec) * bvec - np.dot(direction, bvec_perp) * bvec_perp ) new_direction /= np.linalg.norm(new_direction) return last_point_before_bounce, new_direction @@ -190,9 +183,7 @@ def single_bounce(self, direction, p, step_size=5): percent_random = 0.05 new_direction = ( 1 - percent_random - ) * new_direction + percent_random * np.array( - [np.sin(theta), np.cos(theta)] - ) + ) * new_direction + percent_random * np.array([np.sin(theta), np.cos(theta)]) return to_points, new_direction @@ -255,9 +246,7 @@ def update_status(self, skip_write=False): try: 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) b0_motor_steps, a0_motor_steps, b1_motor_steps, a1_motor_steps = map( float, motor_position_str[len("MPos:") :].split(",") @@ -292,24 +281,19 @@ def wait_while_moving(self): return time.sleep(0.01) - def move_to( - self, points - ): # takes in a list of points equal to length of map + def move_to(self, points): # takes in a list of points equal to length of map gcode_move = ["G0"] for c in points: motors = self.channel_to_motor_map[c] a_motor_steps, b_motor_steps = self.dynamics.to_steps(points[c]) gcode_move += [ - "%s%0.2f %s%0.2f" - % (motors[0], b_motor_steps, motors[1], a_motor_steps) + "%s%0.2f %s%0.2f" % (motors[0], b_motor_steps, motors[1], a_motor_steps) ] cmd = " ".join(gcode_move) time.sleep(0.01) self.s.write((cmd + "\n").encode()) # Send g-code block to grbl time.sleep(0.01) - grbl_out = ( - self.s.readline() - ) # Wait for grbl response with carriage return + grbl_out = self.s.readline() # Wait for grbl response with carriage return time.sleep(0.01) # print("MOVE TO RESPONSE", grbl_out) diff --git a/spf/grbl/mm_per_step.py b/spf/grbl/mm_per_step.py index dd06c0cb..72a2ea58 100644 --- a/spf/grbl/mm_per_step.py +++ b/spf/grbl/mm_per_step.py @@ -7,9 +7,7 @@ steps_per_revolution = 360 / degrees_per_step micro_stepping = 16 steps_per_mm = ( - micro_stepping - * steps_per_revolution - / (teeth_per_revolution * spacing_per_tooth) + micro_stepping * steps_per_revolution / (teeth_per_revolution * spacing_per_tooth) ) diff --git a/spf/grbl_sdr_collect.py b/spf/grbl_sdr_collect.py index 590545c2..92204ac1 100644 --- a/spf/grbl_sdr_collect.py +++ b/spf/grbl_sdr_collect.py @@ -154,9 +154,7 @@ def bounce_grbl(gm): 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"] diff --git a/spf/model_training_and_inference/01_generate_data.py b/spf/model_training_and_inference/01_generate_data.py index f6581bf5..5a6b9dd8 100644 --- a/spf/model_training_and_inference/01_generate_data.py +++ b/spf/model_training_and_inference/01_generate_data.py @@ -15,9 +15,7 @@ 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("--signal-frequency", type=float, required=False, default=100e3) parser.add_argument( "--sampling-frequency", type=float, required=False, default=10e6 ) @@ -29,12 +27,8 @@ 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-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 ) @@ -52,20 +46,12 @@ 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("--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("--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" @@ -74,9 +60,7 @@ 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="+" - ) + parser.add_argument("--fixed-detector", type=float, required=False, nargs="+") args = parser.parse_args() diff --git a/spf/model_training_and_inference/12_task2_model_training.py b/spf/model_training_and_inference/12_task2_model_training.py index 3bc921ea..0ac68762 100644 --- a/spf/model_training_and_inference/12_task2_model_training.py +++ b/spf/model_training_and_inference/12_task2_model_training.py @@ -35,10 +35,7 @@ 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 - ) + 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() @@ -95,10 +92,7 @@ def model_forward(d_model, data, args, train_test_label, update, plot=True): 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 - ): + if preds["transformer_pred"].mean(axis=1).var(axis=0).mean().item() < 1e-13: d_model["dead"] = True else: d_model["dead"] = False @@ -114,9 +108,7 @@ def model_forward(d_model, data, args, train_test_label, update, plot=True): _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"] - ) + 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"]) @@ -145,9 +137,7 @@ def model_forward(d_model, data, args, train_test_label, update, plot=True): 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]) - ) + 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]], @@ -235,9 +225,7 @@ def model_to_losses(running_loss, mean_chunk): [ l[k] for l in running_loss[ - idx - * mean_chunk : (idx + 1) - * mean_chunk + idx * mean_chunk : (idx + 1) * mean_chunk ] ] ) @@ -333,9 +321,7 @@ def plot_loss( 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[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") @@ -345,9 +331,7 @@ def plot_loss( 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"] - ) + 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: @@ -362,9 +346,7 @@ def plot_loss( 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("--embedding-warmup", type=int, required=False, default=0) parser.add_argument( "--snapshots-per-sample", type=int, @@ -373,21 +355,15 @@ def plot_loss( 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("--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("--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 ) @@ -405,9 +381,7 @@ def plot_loss( ) 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("--lr-transformer", type=float, required=False, default=0.00001) parser.add_argument("--plot", type=bool, required=False, default=False) parser.add_argument( "--transformer-input", @@ -467,9 +441,7 @@ def plot_loss( # 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) - ) + ds_test = torch.utils.data.Subset(ds, np.arange(train_size, train_size + test_size)) print("init dataloader") trainloader = torch.utils.data.DataLoader( @@ -496,8 +468,7 @@ def plot_loss( continue models.append( { - "name": "%d snapshots (l%d)" - % (snapshots_per_sample, n_layers), + "name": "%d snapshots (l%d)" % (snapshots_per_sample, n_layers), "model": SnapshotNet( snapshots_per_sample, n_layers=n_layers, @@ -613,12 +584,8 @@ def plot_loss( 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), + "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), } @@ -648,10 +615,7 @@ def prep_data(data): update=i, plot=True, ) - if ( - i % args.lr_scheduler_every - == args.lr_scheduler_every - 1 - ): + 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) @@ -691,15 +655,12 @@ def prep_data(data): update=i, plot=idx == 0, ) - running_losses["test"][d_model["name"]].append( - losses - ) + 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 * 0 + labels.mean(axis=[0, 1], keepdim=True), labels, ).item() } @@ -776,10 +737,7 @@ def prep_data(data): ) 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 - ): + 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, diff --git a/spf/model_training_and_inference/13_learn_beamformer.py b/spf/model_training_and_inference/13_learn_beamformer.py index 0483128e..9b564aef 100644 --- a/spf/model_training_and_inference/13_learn_beamformer.py +++ b/spf/model_training_and_inference/13_learn_beamformer.py @@ -56,9 +56,9 @@ def src_pos_from_radial(inputs, outputs): 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] + torch.stack([torch.cos(theta * np.pi), torch.sin(theta * np.pi)], axis=2)[ + ..., 0 + ] * dist + det_pos ) @@ -88,9 +88,7 @@ def model_forward(d_model, inputs, outputs, args, beamformer): 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"].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 @@ -110,9 +108,7 @@ def model_to_losses(running_loss, mean_chunk): [ l[k] for l in running_loss[ - idx - * mean_chunk : (idx + 1) - * mean_chunk + idx * mean_chunk : (idx + 1) * mean_chunk ] ] ) @@ -166,10 +162,7 @@ def model_to_stats_str(running_loss, mean_chunk): % ( k, "\t".join( - [ - "%0.4f" % v.item() - for v in losses[k][-1][cols_for_loss] - ] + ["%0.4f" % v.item() for v in losses[k][-1][cols_for_loss]] ), ) ) @@ -214,9 +207,7 @@ def plot_loss( 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"] - ) + ax[1].plot(xs[2:], losses["beamformer_loss"][2:], label=d_model["name"]) # _mn=np.min(losses['beamformer_loss']) # if _mn= d_in - self.linear_in = ( - nn.Linear(d_in, d_model) if d_model > d_in else nn.Identity() - ) + 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( @@ -197,10 +195,7 @@ def __init__( d_hid=256, d_embed=64, n_layers=1, - n_outputs=4 - + 4 - + 2 - + 1, # 2 means , 2 variances, 2 angles, responsibility + n_outputs=4 + 4 + 2 + 1, # 2 means , 2 variances, 2 angles, responsibility dropout=0.0, ssn_d_hid=64, ssn_n_layers=8, @@ -263,9 +258,7 @@ def forward(self, x): 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 - ), + torch.zeros(batch_size, time_steps, 1).to(snap_shot_embeddings.device), ], dim=2, ).reshape(batch_size, time_steps, 1, self.d_embed) @@ -300,9 +293,7 @@ def forward(self, x): ) return { "transformer_pred": self_attention_output, - "single_snapshot_pred": single_snapshot_output[ - "single_snapshot_pred" - ], + "single_snapshot_pred": single_snapshot_output["single_snapshot_pred"], } @@ -416,9 +407,9 @@ def forward(self, x): 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) + 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] @@ -429,9 +420,7 @@ def forward(self, x): drone_state_and_observation_embeddings[b, :t][ times_where_this_tracked_is_broadcasting ], - x["times"][b, :t][ - times_where_this_tracked_is_broadcasting - ], + x["times"][b, :t][times_where_this_tracked_is_broadcasting], ] ) ) @@ -460,9 +449,7 @@ def forward(self, x): 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 + 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 @@ -509,9 +496,9 @@ def forward(self, x): ], axis=1, ) - trajectory_predictions[ - b, emitter_idx - ] = self.trajectory_prediction_net(trajectory_input)["output"] + trajectory_predictions[b, emitter_idx] = self.trajectory_prediction_net( + trajectory_input + )["output"] return { "single_snapshot_predictions": single_snapshot_predictions, @@ -752,9 +739,7 @@ def forward(self, x): class UNet(nn.Module): - def __init__( - self, in_channels=3, out_channels=1, width=128, init_features=32 - ): + def __init__(self, in_channels=3, out_channels=1, width=128, init_features=32): super(UNet, self).__init__() features = init_features @@ -769,35 +754,25 @@ def __init__( 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 * 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.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.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.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.decoder2 = UNet._block((features * 2) * 2, features * 2, name="dec2") self.upconv1 = nn.ConvTranspose2d( features * 2, features, kernel_size=2, stride=2 ) diff --git a/spf/model_training_and_inference/test.py b/spf/model_training_and_inference/test.py index 4f4f1222..6f97b90f 100644 --- a/spf/model_training_and_inference/test.py +++ b/spf/model_training_and_inference/test.py @@ -183,17 +183,14 @@ def forward(self, x: Tensor) -> Tensor: train_iter = WikiText2(split="train") tokenizer = get_tokenizer("basic_english") -vocab = build_vocab_from_iterator( - map(tokenizer, train_iter), specials=[""] -) +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 + 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))) @@ -285,14 +282,10 @@ def get_batch(source: Tensor, i: int) -> Tuple[Tensor, Tensor]: ntokens = len(vocab) # size of vocabulary emsize = 200 # embedding dimension d_hid = 200 # dimension of the feedforward network model in ``nn.TransformerEncoder`` -nlayers = ( - 2 # number of ``nn.TransformerEncoderLayer`` in ``nn.TransformerEncoder`` -) +nlayers = 2 # number of ``nn.TransformerEncoderLayer`` in ``nn.TransformerEncoder`` nhead = 2 # number of heads in ``nn.MultiheadAttention`` dropout = 0.2 # dropout probability -model = TransformerModel(ntokens, emsize, nhead, d_hid, nlayers, dropout).to( - device -) +model = TransformerModel(ntokens, emsize, nhead, d_hid, nlayers, dropout).to(device) ###################################################################### @@ -393,9 +386,7 @@ def evaluate(model: nn.Module, eval_data: Tensor) -> float: 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 ###################################################################### @@ -406,8 +397,5 @@ 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(f"| End of training | test loss {test_loss:5.2f} | " f"test ppl {test_ppl:8.2f}") print("=" * 89) diff --git a/spf/model_training_and_inference/utils/image_utils.py b/spf/model_training_and_inference/utils/image_utils.py index 1d1cec05..368283a3 100644 --- a/spf/model_training_and_inference/utils/image_utils.py +++ b/spf/model_training_and_inference/utils/image_utils.py @@ -15,9 +15,9 @@ def get_grid(width): # 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) + 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 @@ -52,9 +52,7 @@ def blur10(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) - ) + 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): @@ -74,9 +72,7 @@ def labels_to_source_images(labels, width, img_width=128): ] = 1 label_images = blur5( - label_images.reshape( - b * s, 1, img_width + 2 * offset, img_width + 2 * offset - ) + 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 @@ -84,14 +80,10 @@ def labels_to_source_images(labels, width, img_width=128): return label_images -def radio_to_image( - beam_former_outputs_at_t, theta_at_pos, detector_orientation -): +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_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 = ( ( diff --git a/spf/model_training_and_inference/utils/plot.py b/spf/model_training_and_inference/utils/plot.py index b34c651c..836d7491 100644 --- a/spf/model_training_and_inference/utils/plot.py +++ b/spf/model_training_and_inference/utils/plot.py @@ -46,18 +46,16 @@ def plot_predictions_and_baseline(session, args, step, pred_a, pred_b): ) # 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]) + 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 - ) + ][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]], @@ -121,9 +119,7 @@ def plot_predictions_and_baseline(session, args, step, pred_a, pred_b): 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 - ) + 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"]) @@ -137,9 +133,7 @@ def plot_predictions_and_baseline(session, args, step, pred_a, pred_b): 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 - ) + 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) @@ -266,11 +260,9 @@ def plot_lines(session, steps, output_prefix): [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 - ) + 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]], @@ -302,9 +294,9 @@ def plot_lines(session, steps, output_prefix): 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( + 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] ) @@ -355,9 +347,7 @@ def plot_lines(session, steps, output_prefix): 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 - ) + 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) @@ -368,9 +358,7 @@ def plot_lines(session, steps, output_prefix): # generate the images for the session -def plot_full_session( - session, steps, output_prefix, img_width=128, invert=False -): +def plot_full_session(session, steps, output_prefix, img_width=128, invert=False): width = session["width_at_t"][0][0] # extract the images @@ -388,9 +376,9 @@ def plot_full_session( 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) + 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): @@ -420,20 +408,18 @@ def plot_full_session( [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 - ) + 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( + 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] ) diff --git a/spf/model_training_and_inference/utils/spf_dataset.py b/spf/model_training_and_inference/utils/spf_dataset.py index 4ecce42e..e1204ff5 100644 --- a/spf/model_training_and_inference/utils/spf_dataset.py +++ b/spf/model_training_and_inference/utils/spf_dataset.py @@ -55,23 +55,16 @@ def __init__(self, root_dir, snapshots_in_sample=5): root_dir (string): Directory with all the images. """ self.root_dir = root_dir - self.args = load( - "/".join([self.root_dir, "args.pkl"]), compression="lzma" - ) + 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.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) - ], + ["%s/%s" % (self.root_dir, x) for x in os.listdir(self.root_dir)], ) ) if self.args.sessions != len( @@ -80,9 +73,7 @@ def __init__(self, root_dir, snapshots_in_sample=5): 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 - ) + 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, @@ -181,10 +172,7 @@ def __init__( self.check_file, filter( lambda x: ".npy" in x, - [ - "%s/%s" % (self.root_dir, x) - for x in os.listdir(self.root_dir) - ], + ["%s/%s" % (self.root_dir, x) for x in os.listdir(self.root_dir)], ), ) ) @@ -207,8 +195,7 @@ def __init__( 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 + 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 = {} @@ -250,9 +237,7 @@ def __getitem__(self, idx): :, None ], # TODO multi source "source_positions_at_t": source_positions_at_t, - "source_velocities_at_t": self.zeros[:, :2][ - :, None - ], # TODO calc velocity, + "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) ), @@ -276,12 +261,8 @@ def __getitem__(self, idx): 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["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], @@ -374,18 +355,14 @@ 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] - ) + 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() + 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): @@ -400,9 +377,7 @@ def collate_fn_beamformer(_in): # 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) + / d["signal_matrixs_at_t"].abs().mean(axis=[2, 3], keepdims=True) ).reshape( b, s, -1 ), # normalize the data @@ -449,8 +424,7 @@ def collate_fn_transformer_filter(_in): normalized_pirads_space_theta = torch.cat( [ torch.zeros(b, 1, 1), - (torch.atan2(space_diffs[..., 1], space_diffs[..., 0]))[:, :, None] - / np.pi, + (torch.atan2(space_diffs[..., 1], space_diffs[..., 0]))[:, :, None] / np.pi, ], axis=1, ) @@ -484,14 +458,10 @@ def collate_fn_transformer_filter(_in): dim=3, ), "emitters_broadcasting": d["broadcasting_positions_at_t"], - "emitters_n_broadcasts": d["broadcasting_positions_at_t"].cumsum( - axis=1 - ), + "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) - ) + 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( @@ -550,8 +520,7 @@ def collate_fn(_in): space_theta = torch.cat( [ torch.zeros(b, 1, 1), - (torch.atan2(space_diffs[..., 1], space_diffs[..., 0]))[:, :, None] - / np.pi, + (torch.atan2(space_diffs[..., 1], space_diffs[..., 0]))[:, :, None] / np.pi, ], axis=1, ) @@ -594,9 +563,7 @@ def collate_fn(_in): ).float(), "radio_feature": torch.cat( [ - torch.log( - d["beam_former_outputs_at_t"].mean(axis=2, keepdim=True) - ) + 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( diff --git a/spf/model_training_and_inference/utils/spf_generate.py b/spf/model_training_and_inference/utils/spf_generate.py index 31ec78c9..faacbbe8 100644 --- a/spf/model_training_and_inference/utils/spf_generate.py +++ b/spf/model_training_and_inference/utils/spf_generate.py @@ -23,9 +23,7 @@ def _arctan2(x, y): class BoundedPoint: - def __init__( - self, pos=np.ones(2) * 0.5, v=np.zeros(2), delta_time=0.05, width=128 - ): + 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 @@ -53,10 +51,7 @@ def time_to_detector_offset( 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 - ) + y = np.sin(2 * np.pi * orbital_frequency * t + phase_offset) * orbital_height return np.array([1 / 2 + x, 1 / 2 + y]) @@ -135,9 +130,7 @@ def generate_session(args_and_session_idx): (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) - ) + 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)) @@ -150,13 +143,10 @@ def generate_session(args_and_session_idx): # 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 + 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 - ) + 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 @@ -172,10 +162,7 @@ def generate_session(args_and_session_idx): 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_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), @@ -184,9 +171,7 @@ def generate_session(args_and_session_idx): ) ) - whos_broadcasting_at_t = np.random.randint( - 0, n_sources_used, args.time_steps - ) + 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) @@ -219,9 +204,7 @@ def generate_session(args_and_session_idx): # 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 - ) + 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)): @@ -237,9 +220,7 @@ def generate_session(args_and_session_idx): d.add_source( NoiseWrapper( IQSource( - current_source_positions[ - [tdm_source_idx] - ], # x, y position + current_source_positions[[tdm_source_idx]], # x, y position args.carrier_frequency, ), sigma=sigma, @@ -256,10 +237,7 @@ def generate_session(args_and_session_idx): orbital_width=1 / 4, orbital_height=1 / 3, phase_offset=detector_position_phase_offset, - orbital_frequency=(2 / 3) - * args.width - * np.pi - / detector_speed, + orbital_frequency=(2 / 3) * args.width * np.pi / detector_speed, ) * args.width ).astype(int) @@ -275,9 +253,7 @@ def generate_session(args_and_session_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 + 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() @@ -301,8 +277,7 @@ def generate_session(args_and_session_idx): if tdm_source_idx >= 0: diff = ( - current_source_positions[tdm_source_idx] - - detector_position_at_t[t_idx] + 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 diff --git a/spf/rf.py b/spf/rf.py index 16448ba5..4fec4fc1 100644 --- a/spf/rf.py +++ b/spf/rf.py @@ -44,8 +44,7 @@ def __init__(self, pos): def signal(self, sampling_times): return ( - np.cos(2 * np.pi * sampling_times) - + np.sin(2 * np.pi * sampling_times) * 1j + np.cos(2 * np.pi * sampling_times) + np.sin(2 * np.pi * sampling_times) * 1j ) def demod_signal(self, signal, demod_times): @@ -62,8 +61,7 @@ def __init__(self, pos, frequency, phase=0, amplitude=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 + + np.sin(2 * np.pi * sampling_times * self.frequency + self.phase) * 1j ) @@ -160,15 +158,10 @@ def all_receiver_pos(self, with_offset=True): if with_offset: return ( self.position_offset - + ( - rotation_matrix(self.orientation) - @ self.receiver_positions.T - ).T + + (rotation_matrix(self.orientation) @ self.receiver_positions.T).T ) else: - return ( - rotation_matrix(self.orientation) @ self.receiver_positions.T - ).T + return (rotation_matrix(self.orientation) @ self.receiver_positions.T).T def receiver_pos(self, receiver_idx, with_offset=True): if with_offset: @@ -221,9 +214,7 @@ def get_signal_matrix(self, start_time, duration, rx_lo=0): 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] - ) + 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 @@ -240,18 +231,14 @@ def get_signal_matrix(self, start_time, duration, rx_lo=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 - ) + 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) - ) + self.set_receiver_positions(linear_receiver_positions(n_elements, spacing)) @functools.lru_cache(maxsize=1024) @@ -263,9 +250,7 @@ def circular_receiver_positions(n_elements, radius): 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) - ) + self.set_receiver_positions(circular_receiver_positions(n_elements, radius)) @functools.lru_cache(maxsize=1024) @@ -382,10 +367,7 @@ def beamformer( carrier_wavelength = c / carrier_frequency args = ( - 2 - * np.pi - * projection_of_receiver_onto_source_directions - / carrier_wavelength + 2 * np.pi * projection_of_receiver_onto_source_directions / carrier_wavelength ) steering_vectors = np.exp(-1j * args) if calibration is not None: @@ -394,7 +376,5 @@ def beamformer( 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 + steer_dot_signal = np.absolute(phase_adjusted).mean(axis=1) # mean over samples return thetas, steer_dot_signal, steering_vectors diff --git a/spf/sdrpluto/01_phase_sync.py b/spf/sdrpluto/01_phase_sync.py index 04f403e6..390e9a5f 100644 --- a/spf/sdrpluto/01_phase_sync.py +++ b/spf/sdrpluto/01_phase_sync.py @@ -46,9 +46,7 @@ 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 +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 @@ -104,9 +102,7 @@ def sample_phase_offset_rx(iterations=64, calibration=1 + 0j): 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 diff --git a/spf/sdrpluto/02_wifi_direction.py b/spf/sdrpluto/02_wifi_direction.py index 076722d3..8069fcbc 100644 --- a/spf/sdrpluto/02_wifi_direction.py +++ b/spf/sdrpluto/02_wifi_direction.py @@ -47,9 +47,7 @@ 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 +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 @@ -105,9 +103,7 @@ def sample_phase_offset_rx(iterations=64, calibration=1 + 0j): 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 diff --git a/spf/sdrpluto/gather.py b/spf/sdrpluto/gather.py index dbf9d175..88f83f47 100644 --- a/spf/sdrpluto/gather.py +++ b/spf/sdrpluto/gather.py @@ -103,16 +103,14 @@ def setup_rxtx_and_phase_calibration(args): 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) + (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)), + "(fraction of 2pi) %0.4f" % (phase_calibrations.mean() / (2 * np.pi)), ) sdr_rxtx.phase_calibration = phase_calibrations.mean() return sdr_rxtx @@ -168,9 +166,7 @@ def setup_rx_and_tx(args): sdr_emitter.tx_lo = int(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) + 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 # @@ -210,9 +206,9 @@ def setup_rx_and_tx(args): 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) + dists = np.vstack([2 * np.pi - np.abs(cm - angles), np.abs(cm - angles)]).min( + axis=0 + ) _angles = angles[dists < np.percentile(dists, 100.0 - trim)] _cm = np.arctan2(np.sin(_angles).sum(), np.cos(_angles).sum()) % (2 * np.pi) return cm, _cm @@ -222,9 +218,7 @@ def get_avg_phase(signal_matrix, trim=0.0): # signal_matrix=np.vstack(sdr_rx.rx()) # signal_matrix[1]*=np.exp(1j*sdr_rx.phase_calibration) - diffs = (np.angle(signal_matrix[0]) - np.angle(signal_matrix[1])) % ( - 2 * np.pi - ) + diffs = (np.angle(signal_matrix[0]) - np.angle(signal_matrix[1])) % (2 * np.pi) mean, _mean = circular_mean(diffs, trim=50.0) return mean, _mean @@ -240,9 +234,7 @@ def plot_recv_signal(sdr_rx): signal_matrix = np.vstack(sdr_rx.rx()) signal_matrix[1] *= np.exp(1j * sdr_rx.phase_calibration) - beam_thetas, beam_sds, beam_steer = beamformer( - pos, signal_matrix, args.fc - ) + beam_thetas, beam_sds, beam_steer = beamformer(pos, signal_matrix, args.fc) freq = np.fft.fftfreq(t.shape[-1], d=1.0 / sdr_rx.sample_rate) assert t.shape[-1] == rx_n @@ -256,22 +248,16 @@ def plot_recv_signal(sdr_rx): axs[idx][0].set_ylim([-1000, 1000]) sp = np.fft.fft(signal_matrix[idx]) - axs[idx][1].scatter( - freq, np.log(np.abs(sp.real)), s=1 - ) # , freq, sp.imag) + axs[idx][1].scatter(freq, np.log(np.abs(sp.real)), s=1) # , freq, sp.imag) axs[idx][1].set_xlabel("Frequency bin") axs[idx][1].set_ylabel("Power") axs[idx][1].set_ylim([-30, 30]) 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].axvline(x=max_freq, label="max %0.2e" % max_freq, color="red") axs[idx][1].legend() axs[idx][2].clear() - axs[idx][2].scatter( - signal_matrix[idx].real, signal_matrix[idx].imag, s=1 - ) + axs[idx][2].scatter(signal_matrix[idx].real, signal_matrix[idx].imag, s=1) axs[idx][2].set_xlabel("I real(signal)") axs[idx][2].set_ylabel("Q imag(signal)") axs[idx][2].set_title("IQ plot recv (%d)" % idx) @@ -281,9 +267,7 @@ def plot_recv_signal(sdr_rx): axs[idx][0].set_title("Real signal recv (%d)" % idx) axs[idx][1].set_title("Power recv (%d)" % idx) # print("MAXFREQ",freq[np.abs(np.argmax(sp.real))]) - diff = (np.angle(signal_matrix[0]) - np.angle(signal_matrix[1])) % ( - 2 * np.pi - ) + diff = (np.angle(signal_matrix[0]) - np.angle(signal_matrix[1])) % (2 * np.pi) axs[0][3].clear() axs[0][3].scatter(t, diff, s=1) mean, _mean = circular_mean(diff) diff --git a/spf/sdrpluto/phase_cal/01_phase_cal.py b/spf/sdrpluto/phase_cal/01_phase_cal.py index 621f10b7..ef4668d7 100644 --- a/spf/sdrpluto/phase_cal/01_phase_cal.py +++ b/spf/sdrpluto/phase_cal/01_phase_cal.py @@ -8,9 +8,7 @@ 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("--ip", type=str, help="target Pluto IP address", required=True) parser.add_argument( "--fi", type=int, help="Intermediate frequency", required=False, default=5e4 ) diff --git a/spf/sdrpluto/sdr.py b/spf/sdrpluto/sdr.py index 2fe5eb16..c96a8421 100644 --- a/spf/sdrpluto/sdr.py +++ b/spf/sdrpluto/sdr.py @@ -19,8 +19,7 @@ def __init__(self, pos): def signal(self, sampling_times): return ( - np.cos(2 * np.pi * sampling_times) - + np.sin(2 * np.pi * sampling_times) * 1j + np.cos(2 * np.pi * sampling_times) + np.sin(2 * np.pi * sampling_times) * 1j ) def demod_signal(self, signal, demod_times): @@ -36,8 +35,7 @@ def __init__(self, pos, frequency, 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 + + np.sin(2 * np.pi * sampling_times * self.frequency + self.phase) * 1j ) @@ -71,8 +69,7 @@ def signal(self, sampling_times): def demod_signal(self, signal, demod_times): return ( - self.lo_in_phase(demod_times) - + self.lo_out_of_phase(demod_times) * 1j + self.lo_in_phase(demod_times) + self.lo_out_of_phase(demod_times) * 1j ) * signal @@ -147,8 +144,7 @@ def beamformer( source_vector, receiver.pos ) projections.append( - projection_of_receiver_onto_source_direction - / carrier_wavelength + projection_of_receiver_onto_source_direction / carrier_wavelength ) arg = ( 2 @@ -158,18 +154,14 @@ def beamformer( ) 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 - ) + 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] - ) + 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 @@ -207,18 +199,14 @@ def plot_space(ax, d, wavelength=1): 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)) - ) + ax.set_title("Space diagram (%s)" % ",".join(map(lambda x: "%0.2f" % x, thetas))) 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]) - ) + self.add_receiver(Receiver([spacing * (idx - (n_elements - 1) / 2), 0])) ula_d = ULADetector(300, 10, 1) diff --git a/spf/sdrpluto/test_emitter_recv/03_only_emit.py b/spf/sdrpluto/test_emitter_recv/03_only_emit.py index 253347f1..21fd0317 100644 --- a/spf/sdrpluto/test_emitter_recv/03_only_emit.py +++ b/spf/sdrpluto/test_emitter_recv/03_only_emit.py @@ -7,9 +7,7 @@ import argparse parser = argparse.ArgumentParser() -parser.add_argument( - "--ip", type=str, help="target Pluto IP address", required=True -) +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 ) diff --git a/spf/sdrpluto/test_emitter_recv/04_plot_signal.py b/spf/sdrpluto/test_emitter_recv/04_plot_signal.py index c5d42972..a566fd9a 100644 --- a/spf/sdrpluto/test_emitter_recv/04_plot_signal.py +++ b/spf/sdrpluto/test_emitter_recv/04_plot_signal.py @@ -9,9 +9,7 @@ # from sdr import * parser = argparse.ArgumentParser() -parser.add_argument( - "--ip", type=str, help="target Pluto IP address", required=True -) +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 ) @@ -71,9 +69,7 @@ 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].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() diff --git a/spf/sdrpluto/test_single_phase/02_wifi_direction.py b/spf/sdrpluto/test_single_phase/02_wifi_direction.py index efa327c1..8a3e1176 100644 --- a/spf/sdrpluto/test_single_phase/02_wifi_direction.py +++ b/spf/sdrpluto/test_single_phase/02_wifi_direction.py @@ -8,9 +8,7 @@ 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("--ip", type=str, help="target Pluto IP address", required=True) parser.add_argument( "--fi", type=int, help="Intermediate frequency", required=False, default=1e5 )