Question: I'm writing a code and I have a function that calculates the values that are not fulfilling a condition with the values that are fulfilling the condition, but I'm having a lot of trouble with managing the shape of the arrays. I have a similar function, but with other logical structure that does this (MWE for the function that works) import numpy as np def f_est(f, fd, mask): """ This function returns the estimated value given a function f and its derivative fd for the points selected by the mask, estimating forward. Parameters: - f: Array of function values. - fd: Array of derivative values of the function. - mask: Boolean mask to select points in f and fd. """ h = 0.000001 # Find the last index that satisfies the mask last_index = np.max(np.where(mask)[0]) # Create shifted masks (inspired by f_cal), but centered on the last index mask_current = mask[:last_index + 1] # Mask for current positions up to the last index mask_prev = mask[:last_index] # Mask for previous positions (for fd_prev_slice) mask_prev2 = mask[:last_index - 1] # Mask for previous positions (for fd_prev2_slice) # Apply masks to f and fd (with shifts), centered on the last index f_slice = f[:last_index + 1][mask_current] # Note: adjusted to align with mask_current fd_slice = fd[:last_index + 1][mask_current] fd_prev_slice = fd[:last_index][mask_prev] fd_prev2_slice = fd[:last_index - 1][mask_prev2] # Perform the calculations with consistent slices, estimating forward # Use the last value of f_slice, fd_slice, fd_prev_slice, and fd_prev2_slice for estimation last_f = f_slice[-1] last_fd = fd_slice[-1] last_fd_prev = fd_prev_slice[-1] if len(fd_prev_slice) > 0 else 0 last_fd_prev2 = fd_prev2_slice[-1] if len(fd_prev2_slice) > 0 else 0 estimated_next_value = ( last_f + h * last_fd + 1 / 2 * (h * last_fd - h * last_fd_prev) + 5 / 12 * ((h * last_fd - h * last_fd_prev) - (h * last_fd_prev - h * last_fd_prev2)) ) return estimated_next_value f = np.array([1, 2, 3, 4, 5, 6, 7]) fd = f mask = np.array([True, True, True, False, False, False, False]) print("Original Array:", f) print("Length of Original Array:", len(f)) print("Masked Array:", f[~mask]) print("Length of Masked Array:", len(f[~mask])) f[~mask] = f_est(f, fd, mask) print("Final Array:", f) print("Length of Final Array:", len(f))  But with this function (MWE that doesn't work): import numpy as np def f_cal(var, dvar, mask): """ Calculates next value using trapezoidal method with masks, estimating forward. Parameters: var: array of current values dvar: array of derivatives mask: boolean array indicating which positions to calculate """ h = 0.0000001 # Encontrar el último índice que verifica la máscara last_index = np.max(np.where(mask)[0]) # Crear máscaras para posiciones actuales y la siguiente mask_current = mask[:last_index] mask_next = mask[1:last_index+1] # Marcar como True el índice siguiente # Ajustar los arreglos para alinear con las máscaras var_current = var[:last_index+1][mask_current] dvar_current = dvar[:last_index+1][mask_current] dvar_next = dvar[:last_index+2][mask_next][:1] # Solo el valor siguiente # Calculate using trapezoidal method with masks, estimating forward result = var_current + h * dvar_next - 1/2*(h*dvar_next-h*dvar_current[-1]) return result f = np.array([1, 2, 3, 4, 5,6,7]) fd = f mask = np.array([True, True, True, False, False, False, False]) print("Original Array:", f) print("Length of Original Array:", len(f)) print("Masked Array:", f[~mask]) print("Length of Masked Array:", len(f[~mask])) f[~mask] = f_cal(f, fd, mask) print("Final Array:", f) print("Length of Final Array:", len(f))  I'm having a lot of trouble keeping the length of the array to match the number of elements that are not satisfying the condition

Answer: I believe you're trying to use the Taylor series to approximate the function at a position beyond the last known value. The function values are stored in the var array, and the first derivatives are in the dvar array. It seems like you're using mask to identify the last known function value, and h represents the argument step. From your code, it looks like you're using terms up to the second derivative, which you approximate as the difference between adjacent first derivatives divided by the argument step. This can be expressed as:  f(x+h) = f(x) + h⋅f′(x) + ½⋅h²⋅f″(x) + … ≈ ≈ f(x) + h⋅f′(x) + ½⋅h²⋅(f′(x+h)-f′(x))/h  At least, this resonates with the expression you return as the result (which I tweaked slightly):  var_current + h*dvar_next - 1/2*(h*dvar_next-h*dvar_current)  We can transform the above line step by step as follows:  var_current + h*dvar_next - 1/2*h*dvar_next + 1/2*h*dvar_current == var_current + 1/2*h*dvar_next + 1/2*h*dvar_current == var_current + h*dvar_current + 1/2*h*dvar_next - 1/2*h*dvar_current == var_current + h*dvar_current + 1/2*h*(dvar_next - dvar_current) == var_current + h*dvar_current + 1/2*(h**2)*(dvar_next - dvar_current)/h  The last line reminded me of the Taylor series, which is why I suspect that the mention of the Trapezoidal method in your function description might be a mistake. If that's the case, the function could look like this: def f_cal(var, dvar, mask, h=1e-7): index = np.asarray(mask).nonzero()[0] assert index.size > 0, 'Nowhere to start from' assert index[-1] < len(mask)-1, 'Nowhere to extrapolate to' pos = index[-1] f = var[pos] df = dvar[pos] df_next = dvar[pos+1] return f + h*df + h*(df_next - df)/2  However, I might have misunderstood your intent. For example, if your reference to the trapezoidal method is deliberate, or you're using a different convention, you may need to clarify your question.  P.S. Answer to the original question about mismatching lengths when using a mask: you seem to have missed handling array dimensions in some places. Here's an alternative implementation to compare with (note that the changes are made based on the assumption that you approximate the function using Taylor series): def f_cal(var, dvar, mask): h = 0.0000001 last_index = np.max(np.where(mask)[0]) mask_current = mask[:last_index+1] mask_next = mask[:last_index+2] var_current = var[:last_index+1][mask_current][-1] dvar_current = dvar[:last_index+1][mask_current][-1] dvar_next = dvar[:last_index+2][mask_next][-1] return var_current + h*dvar_next - 1/2*h*(dvar_next-dvar_current)