Merge branch 'new_intensity'

Author: ahillsley

blinx/estimate.py

@@ -12,7 +12,7 @@
 # FIXME: post_process should be renamed and find a new home
 from .post_process import post_process as find_most_likely_y
 from .trace_model import get_trace_log_likelihood
-from .utils import find_local_maxima
+from .utils import find_maximum
 
 
 def estimate_y(traces, max_y, parameter_ranges=None, hyper_parameters=None):
@@ -246,9 +246,9 @@ def get_initial_parameter_guesses(traces, y, parameter_ranges, hyper_parameters)
         )
 
         # find locations where parameters maximize log likelihoods
-        min_indices = find_local_maxima(trace_log_likelihoods, num_guesses)
+        min_index = find_maximum(trace_log_likelihoods)
 
-        guesses.append(parameters[min_indices])
+        guesses.append(parameters[min_index])
 
     # all guesses are stored in 'guesses', the following stacks them together
     # as if we vmap'ed over traces:

blinx/trace_model.py

@@ -105,6 +105,14 @@ def single_optimal_trace(trace, y, parameters, hyper_parameters):
 
     """
 
+    r_e = parameters.r_e
+    r_bg = parameters.r_bg
+    mu_ro = parameters.mu_ro
+    sigma_ro = parameters.sigma_ro
+    gain = parameters.gain
+    p_on = parameters.p_on
+    p_off = parameters.p_off
+
     zs = jnp.arange(0, y + 1)
 
     # Discretize the trace into bins
@@ -116,20 +124,12 @@ def single_optimal_trace(trace, y, parameters, hyper_parameters):
     x_left = (trace // bin_width) * bin_width
     x_right = x_left + bin_width
 
-    p_transition = create_transition_matrix(y, parameters.p_on, parameters.p_off)
+    p_transition = create_transition_matrix(y, p_on, p_off)
     p_initial = get_steady_state(p_transition)
     p_measurement = jax.vmap(
         p_x_given_z,
-        in_axes=(None, None, 0, None, None, None, None),
-    )(
-        x_left,
-        x_right,
-        zs,
-        parameters.mu,
-        parameters.mu_bg,
-        parameters.sigma,
-        hyper_parameters,
-    )
+        in_axes=(None, None, 0, None, None, None, None, None, None),
+    )(x_left, x_right, zs, r_e, r_bg, mu_ro, sigma_ro, gain, hyper_parameters)
 
     return get_optimal_states(p_measurement, p_initial, p_transition)
 

blinx/utils.py

@@ -3,46 +3,18 @@
 import scipy.signal
 
 
-def find_local_maxima(matrix, num_maxima=None):
-    # convert to numpy array
-    matrix = np.asarray(matrix)
+def find_maximum(matrix):
+    temp_matrix = np.array(matrix)
 
-    if num_maxima is None:
-        num_maxima = matrix.size
+    # argmax will return index of any nans
+    # so replace nans with -inf
 
-    # pad matrix with -inf
-    padded = np.ones(tuple(s + 2 for s in matrix.shape), dtype=matrix.dtype)
-    padded *= -np.inf
-    slices = tuple(slice(1, s + 1) for s in matrix.shape)
-    padded[slices] = matrix
+    mask = np.isnan(temp_matrix)
 
-    padded_indices = scipy.signal.argrelmax(np.asarray(padded), mode="wrap")
-    # indices into original matrix without padding
-    indices = tuple(i - 1 for i in padded_indices)
+    temp_matrix[mask] = -np.inf
 
-    # set all non-maxima to -inf
-    maxima = np.ones_like(matrix)
-    maxima *= -np.inf
-    maxima[indices] = matrix[indices]
+    index = np.argmax(temp_matrix)
 
-    # get all maximum values, sorted
-    values, indices = np.unique(maxima, return_index=True)
+    index = np.unravel_index(index, temp_matrix.shape)
 
-    assert values[0] == -np.inf
-
-    # first index should point to -np.inf, drop it
-    indices = indices[1:]
-
-    # values, indices = np.unique(np.isfinite(maxima), return_index=True)
-
-    # retain only last num_maxima values
-    if len(indices) > num_maxima:
-        indices = indices[-num_maxima:]
-
-    # convert back to non-flattened indices
-    indices = np.unravel_index(indices, matrix.shape)
-
-    # convert to jax array
-    indices = tuple(jnp.array(i) for i in indices)
-
-    return indices
+    return tuple(jnp.expand_dims(jnp.array(i), axis=-1) for i in index)