Remove slicing that offended mypy

indica-mcf · Oct 28, 2022 · 363ab66 · 363ab66
1 parent 5dab824
commit 363ab66
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/indica/operators/extrapolate_impurity_density.py b/indica/operators/extrapolate_impurity_density.py
@@ -837,7 +837,9 @@ def optimize_perturbation(
             orig_bolometry = orig_bolometry_data
 
         extrapolated_smooth_data_mean = np.mean(
-            extrapolated_smooth_data.loc[threshold_rho[0] :, :, :]
+            extrapolated_smooth_data.sel(
+                rho_poloidal=slice(threshold_rho.isel(t=0), None)
+            )
         )
 
         def objective_func(objective_array: Sequence, time: float):