Fix tests after merge

Author: cwindolf

src/dartsort/localize/localize_torch.py

@@ -69,8 +69,7 @@ def localize_amplitude_vectors(
     assert channel_index.shape == (n_channels_tot, c)
     assert main_channels.shape == (n_spikes,)
     # we'll return numpy if user sent numpy
-     is_numpy = not torch.is_tensor(amplitude_vectors)
-
+    is_numpy = not torch.is_tensor(amplitude_vectors)
 
     # handle channel subsetting
     if radius is not None or n_channels_subset is not None:
@@ -114,26 +113,22 @@ def localize_amplitude_vectors(
 
     if model == "com":
         z_abs_com = zcom + geom[main_channels, 1]
-        nancom = torch.full_like(xcom, torch.nan)
-        return dict(
-            x=xcom, y=nancom, z_rel=zcom, z_abs=z_abs_com, alpha=nancom
-        )
+        return dict(x=xcom, z_rel=zcom, z_abs=z_abs_com)
 
     # normalized PTP vectors
     # this helps to keep the objective in a similar range, so we can use
     # fixed constants in regularizers like the log barrier
     max_amplitudes = torch.max(amplitude_vectors, dim=1).values
     normalized_amp_vecs = amplitude_vectors / max_amplitudes[:, None]
-    
+
     # -- torch optimize
+    if levenberg_marquardt_kwargs is None:
+        levenberg_marquardt_kwargs = {}
+
     # initialize with center of mass
     locs = torch.column_stack((xcom, torch.full_like(xcom, y0), zcom))
 
-    
     if model == "pointsource":
-
-        if levenberg_marquardt_kwargs is None:
-            levenberg_marquardt_kwargs = {}
         locs, i = batched_levenberg_marquardt(
             locs,
             vmap_point_source_grad_and_mse,
@@ -149,52 +144,53 @@ def localize_amplitude_vectors(
             amplitude_vectors, in_probe_mask, x, y, z_rel, local_geoms
         )
         z_abs = z_rel + geom[main_channels, 1]
-        
+
+        results = dict(x=x, y=y, z_rel=z_rel, z_abs=z_abs, alpha=alpha)
         if is_numpy:
-            x = x.numpy(force=True)
-            y = y.numpy(force=True)
-            z_rel = z_rel.numpy(force=True)
-            z_abs = z_abs.numpy(force=True)
-            alpha = alpha.numpy(force=True)
-        return dict(x=x, y=y, z_rel=z_rel, z_abs=z_abs, alpha=alpha)
-    
-    if model == "dipole":
-        if levenberg_marquardt_kwargs is None:
-            levenberg_marquardt_kwargs = {}
+            results = {k: v.numpy(force=True) for k, v in results.items()}
+        return results
+
+    elif model == "dipole":
         locs, i = batched_levenberg_marquardt(
             locs,
             vmap_dipole_grad_and_mse,
             vmap_dipole_hessian,
             extra_args=(normalized_amp_vecs, local_geoms),
             **levenberg_marquardt_kwargs,
         )
-        
+
         x, y0, z_rel = locs.T
         y = F.softplus(y0)
         projected_dist = vmap_dipole_find_projection_distance(
             normalized_amp_vecs, x, y, z_rel, local_geoms
-        )        
-        
+        )
+
         # if projected_dist>th_dipole_proj_dist: return the loc values from pointsource
 
-        pointsource_spikes = torch.nonzero(projected_dist>th_dipole_proj_dist, as_tuple=True)
-        
+        pointsource_spikes = torch.nonzero(
+            projected_dist > th_dipole_proj_dist, as_tuple=True
+        )
+
         locs_pointsource_spikes, i = batched_levenberg_marquardt(
             locs[pointsource_spikes],
             vmap_point_source_grad_and_mse,
             vmap_point_source_hessian,
-            extra_args=(normalized_amp_vecs[pointsource_spikes], in_probe_mask, local_geoms[pointsource_spikes]),
+            extra_args=(
+                normalized_amp_vecs[pointsource_spikes],
+                in_probe_mask,
+                local_geoms[pointsource_spikes],
+            ),
             **levenberg_marquardt_kwargs,
         )
         x_pointsource_spikes, y0_pointsource_spikes, z_rel_pointsource_spikes = locs.T
         y_pointsource_spikes = F.softplus(y0_pointsource_spikes)
-        
+
         x[pointsource_spikes] = x_pointsource_spikes
         y[pointsource_spikes] = y_pointsource_spikes
         z_rel[pointsource_spikes] = z_rel_pointsource_spikes
-        
+
         z_abs = z_rel + geom[main_channels, 1]
-        
+
         if is_numpy:
             x = x.numpy(force=True)
             y = y.numpy(force=True)
@@ -204,9 +200,14 @@ def localize_amplitude_vectors(
 
         return dict(x=x, y=y, z_rel=z_rel, z_abs=z_abs, alpha=projected_dist)
 
+    else:
+        assert False
+
 
 # -- point source / dipole model library functions
-def point_source_amplitude_at(x, y, z, local_geom):
+
+
+def point_source_amplitude_at(x, y, z, alpha, local_geom):
     """Point source model predicted amplitude at local_geom given location"""
     dxs = torch.square(x - local_geom[:, 0])
     dzs = torch.square(z - local_geom[:, 1])
@@ -223,25 +224,8 @@ def point_source_find_alpha(amp_vec, channel_mask, x, y, z, local_geoms):
     )
     return alpha
 
-def dipole_find_projection_distance(normalized_amp_vec, x, y, z, local_geom):
-    """We can solve for the brightness (alpha) of the source in closed form given x,y,z"""
-    
-    dxs = x - local_geom[:, 0]
-    dzs = z - local_geom[:, 1]
-    dys = y
-    duv = torch.tensor([dxs, dys, dzs])
-    X = duv / torch.pow(torch.sum(torch.square(duv)), 3/2)
-    beta = torch.linalg.solve(torch.matmul(X.T, X), torch.matmul(X.T, normalized_amp_vec))
-    beta /= torch.sqrt(torch.square(beta).sum())
-    dipole_planar_direction = torch.sqrt(np.torch(beta[[0, 2]]).sum())
-    closest_chan = torch.square(duv).sum(1).argmin()
-    min_duv = duv[closest_chan]
-    val_th = torch.sqrt(torch.square(min_duv).sum())/dipole_planar_direction
-    return val_th
 
-def point_source_mse(
-    loc, amplitude_vector, channel_mask, local_geom, logbarrier=True
-):
+def point_source_mse(loc, amplitude_vector, channel_mask, local_geom, logbarrier=True):
     """Objective in point source model
 
     Arguments
@@ -264,36 +248,53 @@ def point_source_mse(
     x, y0, z = loc
     y = F.softplus(y0)
 
-    alpha = point_source_find_alpha(
-        amplitude_vector, channel_mask, x, y, z, local_geom
-    )
+    alpha = point_source_find_alpha(amplitude_vector, channel_mask, x, y, z, local_geom)
     obj = torch.square(
-        amplitude_vector
-        - point_source_amplitude_at(x, y, z, alpha, local_geom)
+        amplitude_vector - point_source_amplitude_at(x, y, z, alpha, local_geom)
     ).mean()
     if logbarrier:
         obj -= torch.log(10.0 * y) / 10000.0
         # idea for logbarrier on points which run away
         # obj -= torch.log(1000.0 - torch.sqrt(torch.square(x) + torch.square(z))).sum() / 10000.0
     return obj
 
+
+def dipole_find_projection_distance(normalized_amp_vec, x, y, z, local_geom):
+    """We can solve for the brightness (alpha) of the source in closed form given x,y,z"""
+
+    dxs = x - local_geom[:, 0]
+    dzs = z - local_geom[:, 1]
+    dys = y
+    duv = torch.tensor([dxs, dys, dzs])
+    X = duv / torch.pow(torch.sum(torch.square(duv)), 3 / 2)
+    beta = torch.linalg.solve(
+        torch.matmul(X.T, X), torch.matmul(X.T, normalized_amp_vec)
+    )
+    beta /= torch.sqrt(torch.square(beta).sum())
+    dipole_planar_direction = torch.sqrt(np.torch(beta[[0, 2]]).sum())
+    closest_chan = torch.square(duv).sum(1).argmin()
+    min_duv = duv[closest_chan]
+    val_th = torch.sqrt(torch.square(min_duv).sum()) / dipole_planar_direction
+    return val_th
+
+
 def dipole_mse(loc, amplitude_vector, local_geom, logbarrier=True):
     """Dipole model predicted amplitude at local_geom given location"""
-    
+
     x, y0, z = loc
     y = F.softplus(y0)
 
     dxs = x - local_geom[:, 0]
     dzs = z - local_geom[:, 1]
     dys = y
-    
+
     duv = torch.tensor([dxs, dys, dzs])
 
-    X = duv / torch.pow(torch.sum(torch.square(duv)), 3/2)
-    
+    X = duv / torch.pow(torch.sum(torch.square(duv)), 3 / 2)
+
     beta = torch.linalg.solve(torch.matmul(X.T, X), torch.matmul(X.T, (ptp / maxptp)))
     qtq = torch.matmul(X, beta)
-    
+
     obj = torch.square(ptp / maxptp - qtq).mean()
     if logbarrier:
         obj -= torch.log(10.0 * y) / 10000.0