1
- use self :: interpolate:: Interpolate ;
1
+ use self :: interpolate:: { higher_index , lower_index , Interpolate } ;
2
2
use super :: sort:: get_many_from_sorted_mut_unchecked;
3
3
use indexmap:: { IndexMap , IndexSet } ;
4
4
use ndarray:: prelude:: * ;
@@ -204,57 +204,51 @@ where
204
204
205
205
let axis_len = self . len_of ( axis) ;
206
206
if axis_len == 0 {
207
- None
208
- } else {
209
- let mut deduped_qs: Vec < N64 > = qs. to_vec ( ) ;
210
- deduped_qs. sort_by ( |a, b| a. partial_cmp ( b) . unwrap ( ) ) ;
211
- deduped_qs. dedup ( ) ;
207
+ return None ;
208
+ }
209
+
210
+ let mut deduped_qs: Vec < N64 > = qs. to_vec ( ) ;
211
+ deduped_qs. sort_by ( |a, b| a. partial_cmp ( b) . unwrap ( ) ) ;
212
+ deduped_qs. dedup ( ) ;
212
213
213
- // IndexSet preserves insertion order:
214
- // - indexes will stay sorted;
215
- // - we avoid index duplication.
216
- let mut searched_indexes = IndexSet :: new ( ) ;
217
- for q in deduped_qs. iter ( ) {
218
- if I :: needs_lower ( * q, axis_len) {
219
- searched_indexes. insert ( I :: lower_index ( * q, axis_len) ) ;
220
- }
221
- if I :: needs_higher ( * q, axis_len) {
222
- searched_indexes. insert ( I :: higher_index ( * q, axis_len) ) ;
223
- }
214
+ // IndexSet preserves insertion order:
215
+ // - indexes will stay sorted;
216
+ // - we avoid index duplication.
217
+ let mut searched_indexes = IndexSet :: new ( ) ;
218
+ for q in deduped_qs. iter ( ) {
219
+ if I :: needs_lower ( * q, axis_len) {
220
+ searched_indexes. insert ( lower_index ( * q, axis_len) ) ;
221
+ }
222
+ if I :: needs_higher ( * q, axis_len) {
223
+ searched_indexes. insert ( higher_index ( * q, axis_len) ) ;
224
224
}
225
- let searched_indexes: Vec < usize > = searched_indexes. into_iter ( ) . collect ( ) ;
225
+ }
226
+ let searched_indexes: Vec < usize > = searched_indexes. into_iter ( ) . collect ( ) ;
226
227
227
- // Retrieve the values corresponding to each index for each slice along the specified axis
228
- let values = self . map_axis_mut ( axis, |mut x| {
229
- get_many_from_sorted_mut_unchecked ( & mut x, & searched_indexes)
230
- } ) ;
228
+ // Retrieve the values corresponding to each index for each slice along the specified axis
229
+ let values = self . map_axis_mut (
230
+ axis,
231
+ |mut x| get_many_from_sorted_mut_unchecked ( & mut x, & searched_indexes)
232
+ ) ;
231
233
232
- // Combine the retrieved values according to specified interpolation strategy to
233
- // get the desired quantiles
234
- let mut results = IndexMap :: new ( ) ;
235
- for q in qs {
236
- let result = I :: interpolate (
237
- match I :: needs_lower ( * q, axis_len) {
238
- true => {
239
- let lower_index = & I :: lower_index ( * q, axis_len) ;
240
- Some ( values. map ( |x| x. get ( lower_index) . unwrap ( ) . clone ( ) ) )
241
- }
242
- false => None ,
243
- } ,
244
- match I :: needs_higher ( * q, axis_len) {
245
- true => {
246
- let higher_index = & I :: higher_index ( * q, axis_len) ;
247
- Some ( values. map ( |x| x. get ( higher_index) . unwrap ( ) . clone ( ) ) )
248
- }
249
- false => None ,
250
- } ,
251
- * q,
252
- axis_len,
253
- ) ;
254
- results. insert ( * q, result) ;
255
- }
256
- Some ( results)
234
+ // Combine the retrieved values according to specified interpolation strategy to
235
+ // get the desired quantiles
236
+ let mut results = IndexMap :: new ( ) ;
237
+ for q in qs {
238
+ let lower = if I :: needs_lower ( * q, axis_len) {
239
+ Some ( values. map ( |x| x[ & lower_index ( * q, axis_len) ] . clone ( ) ) )
240
+ } else {
241
+ None
242
+ } ;
243
+ let higher = if I :: needs_higher ( * q, axis_len) {
244
+ Some ( values. map ( |x| x[ & higher_index ( * q, axis_len) ] . clone ( ) ) )
245
+ } else {
246
+ None
247
+ } ;
248
+ let interpolated = I :: interpolate ( lower, higher, * q, axis_len) ;
249
+ results. insert ( * q, interpolated) ;
257
250
}
251
+ Some ( results)
258
252
}
259
253
260
254
fn quantile_axis_mut < I > ( & mut self , axis : Axis , q : N64 ) -> Option < Array < A , D :: Smaller > >
@@ -276,24 +270,23 @@ where
276
270
S : DataMut ,
277
271
I : Interpolate < A :: NotNan > ,
278
272
{
279
- if self . len_of ( axis) > 0 {
280
- Some ( self . map_axis_mut ( axis, |lane| {
281
- let mut not_nan = A :: remove_nan_mut ( lane) ;
282
- A :: from_not_nan_opt ( if not_nan. is_empty ( ) {
283
- None
284
- } else {
285
- Some (
286
- not_nan
287
- . quantile_axis_mut :: < I > ( Axis ( 0 ) , q)
288
- . unwrap ( )
289
- . into_raw_vec ( )
290
- . remove ( 0 ) ,
291
- )
292
- } )
293
- } ) )
294
- } else {
295
- None
273
+ if self . len_of ( axis) == 0 {
274
+ return None ;
296
275
}
276
+ let quantile = self . map_axis_mut ( axis, |lane| {
277
+ let mut not_nan = A :: remove_nan_mut ( lane) ;
278
+ A :: from_not_nan_opt ( if not_nan. is_empty ( ) {
279
+ None
280
+ } else {
281
+ Some (
282
+ not_nan
283
+ . quantile_axis_mut :: < I > ( Axis ( 0 ) , q)
284
+ . unwrap ( )
285
+ . into_scalar ( ) ,
286
+ )
287
+ } )
288
+ } ) ;
289
+ Some ( quantile)
297
290
}
298
291
}
299
292
0 commit comments