1use comnum_core::signal::{KNOWN_BIT_LLR, Symbols};
72use comnum_core::traits::Detect;
73use comnum_core::{Error, Result, ensure, ensure_shape};
74use comnum_mapping::constellation::Constellation;
75use comnum_math::C64;
76use comnum_math::Jacobian;
77use comnum_math::linalg::{CMat, hermitian, identity, inverse_hpd, matvec, qr};
78use ndarray::{Array2, ArrayView2};
79
80pub const ML_MAX_HYPOTHESES: usize = 1 << 22;
82
83pub(crate) const FLOOR: f64 = 1e-300;
86
87#[derive(Clone, Debug, Default, PartialEq)]
92pub struct Estimates {
93 pub symbols: Vec<C64>,
95 pub noise_var: Vec<f64>,
97}
98
99pub trait MimoDetector {
101 fn constellation(&self) -> &Constellation;
104
105 fn has_soft_output(&self) -> bool {
107 false
108 }
109
110 fn soft_estimate(&self, _h: &ArrayView2<'_, C64>, _y: &[C64], _noise_var: f64) -> Result<Option<Estimates>> {
116 Ok(None)
117 }
118
119 fn detect(&self, h: &ArrayView2<'_, C64>, y: &[C64], noise_var: f64) -> Result<Vec<usize>>;
126
127 fn llr(&self, _h: &ArrayView2<'_, C64>, _y: &[C64], _noise_var: f64) -> Result<Vec<f64>> {
133 Err(Error::Config("this detector has no soft output".into()))
134 }
135
136 fn detect_llr(&self, h: &ArrayView2<'_, C64>, y: &[C64], noise_var: f64) -> Result<(Vec<usize>, Vec<f64>)> {
143 Ok((self.detect(h, y, noise_var)?, self.llr(h, y, noise_var)?))
144 }
145
146 fn detect_block(&self, h: &ArrayView2<'_, C64>, y: &ArrayView2<'_, C64>, noise_var: f64) -> Result<Array2<usize>> {
152 let mut out = Array2::zeros((h.ncols(), y.ncols()));
153 for (t, col) in y.columns().into_iter().enumerate() {
154 let d = self.detect(h, &col.to_vec(), noise_var)?;
155 out.column_mut(t).assign(&ndarray::Array1::from(d));
156 }
157 Ok(out)
158 }
159}
160
161pub(crate) fn check_shape(h: &ArrayView2<'_, C64>, y: &[C64]) -> Result<()> {
163 ensure_shape(y.len() == h.nrows(), || format!("{} received samples for {} receive antennas", y.len(), h.nrows()))
164}
165
166fn check(h: &ArrayView2<'_, C64>, y: &[C64], noise_var: f64) -> Result<()> {
167 check_shape(h, y)?;
168 ensure(noise_var > 0.0 && noise_var.is_finite(), || format!("the noise variance must be positive, got {noise_var}"))
169}
170
171fn tall(h: &ArrayView2<'_, C64>, what: &str) -> Result<()> {
172 ensure(h.nrows() >= h.ncols(), || {
173 format!("{what} needs at least as many receive ({}) as transmit ({}) antennas", h.nrows(), h.ncols())
174 })
175}
176
177fn check_block(h: &ArrayView2<'_, C64>, y: &ArrayView2<'_, C64>, noise_var: f64) -> Result<()> {
179 if y.nrows() != h.nrows() {
180 return Err(Error::Shape(format!("{} received streams for {} receive antennas", y.nrows(), h.nrows())));
181 }
182 check(h, &vec![C64::new(0.0, 0.0); h.nrows()], noise_var)
183}
184
185fn by_column(
187 nt: usize,
188 y: &ArrayView2<'_, C64>,
189 mut detect: impl FnMut(&mut [C64]) -> Result<Vec<usize>>,
190) -> Result<Array2<usize>> {
191 let mut out = Array2::zeros((nt, y.ncols()));
192 let mut col = vec![C64::new(0.0, 0.0); y.nrows()];
193 for (t, c) in y.columns().into_iter().enumerate() {
194 col.iter_mut().zip(c).for_each(|(a, b)| *a = *b);
195 for (o, l) in out.column_mut(t).iter_mut().zip(detect(&mut col)?) {
196 *o = l;
197 }
198 }
199 Ok(out)
200}
201
202#[derive(Clone, Copy, Debug, PartialEq, Eq)]
204pub enum Criterion {
205 Mr,
209 Zf,
211 Mmse,
213}
214
215struct Filter {
217 w: CMat,
218 bias: Vec<f64>,
219 var: Vec<f64>,
220}
221
222fn regularisation(h: &ArrayView2<'_, C64>, criterion: Criterion, es: f64, noise_var: f64) -> Result<f64> {
225 match criterion {
226 Criterion::Mr => Ok(0.0),
227 Criterion::Zf => tall(h, "zero forcing").map(|()| 0.0),
228 Criterion::Mmse => Ok(noise_var / es),
229 }
230}
231
232fn filter(h: &ArrayView2<'_, C64>, criterion: Criterion, es: f64, noise_var: f64) -> Result<Filter> {
233 let n = h.ncols();
234 let lambda = regularisation(h, criterion, es, noise_var)?;
235 if n == 1 {
236 let s = SingleStream::new(h, criterion, es, noise_var, lambda)?;
237 return Ok(Filter { w: h.t().mapv(|v| v.conj() / s.g), bias: vec![s.bias], var: vec![s.var] });
238 }
239 let hh = hermitian(h);
240 if criterion == Criterion::Mr {
241 return matched_filter(&hh, h, es, noise_var);
242 }
243 let g = hh.dot(h) + identity(n).mapv(|x| x * lambda);
244 let ginv = inverse_hpd(&g.view()).map_err(|_| Error::Numerical("the channel is rank deficient".into()))?;
245 let w = ginv.dot(&hh);
246 let (mut bias, mut var) = (Vec::with_capacity(n), Vec::with_capacity(n));
247 for k in 0..n {
248 let (b, v) = unbiased(criterion, noise_var * ginv[[k, k]].re, es);
249 bias.push(b);
250 var.push(v);
251 }
252 Ok(Filter { w, bias, var })
253}
254
255fn matched_filter(hh: &CMat, h: &ArrayView2<'_, C64>, es: f64, noise_var: f64) -> Result<Filter> {
259 let n = h.ncols();
260 let g = hh.dot(h);
261 let mut w = hh.clone();
262 let mut var = Vec::with_capacity(n);
263 for k in 0..n {
264 let gkk = g[[k, k]].re;
265 if !(gkk > 0.0 && gkk.is_finite()) {
266 return Err(Error::Numerical("a stream has a zero channel".into()));
267 }
268 w.row_mut(k).mapv_inplace(|v| v / gkk);
269 let interference: f64 = (0..n).filter(|j| *j != k).map(|j| g[[k, j]].norm_sqr()).sum();
270 var.push((noise_var * gkk + es * interference) / (gkk * gkk));
271 }
272 Ok(Filter { w, bias: vec![1.0; n], var })
273}
274
275fn unbiased(criterion: Criterion, ekk: f64, es: f64) -> (f64, f64) {
279 match criterion {
280 Criterion::Mr | Criterion::Zf => (1.0, ekk),
281 Criterion::Mmse => ((1.0 - ekk / es).max(FLOOR), es * ekk / (es - ekk).max(FLOOR)),
282 }
283}
284
285struct SingleStream {
290 g: f64,
291 bias: f64,
292 var: f64,
293}
294
295impl SingleStream {
296 fn new(
297 h: &ArrayView2<'_, C64>,
298 criterion: Criterion,
299 es: f64,
300 noise_var: f64,
301 lambda: f64,
302 ) -> Result<SingleStream> {
303 let g = h.iter().map(C64::norm_sqr).sum::<f64>() + lambda;
304 if !(g > 0.0 && g.is_finite()) {
305 return Err(Error::Numerical("the channel is rank deficient".into()));
306 }
307 let (bias, var) = unbiased(criterion, noise_var / g, es);
308 Ok(SingleStream { g, bias, var })
309 }
310
311 fn estimate(&self, h: &ArrayView2<'_, C64>, y: &[C64]) -> C64 {
313 h.column(0).iter().zip(y).map(|(hi, yi)| hi.conj() / self.g * yi).sum::<C64>() / self.bias
314 }
315}
316
317#[derive(Clone, Debug, PartialEq)]
319pub struct Linear {
320 pub constellation: Constellation,
322 pub criterion: Criterion,
324}
325
326impl Linear {
327 pub fn mr(constellation: Constellation) -> Linear {
329 Linear { constellation, criterion: Criterion::Mr }
330 }
331
332 pub fn zf(constellation: Constellation) -> Linear {
334 Linear { constellation, criterion: Criterion::Zf }
335 }
336
337 pub fn mmse(constellation: Constellation) -> Linear {
339 Linear { constellation, criterion: Criterion::Mmse }
340 }
341
342 pub fn estimate(&self, h: &ArrayView2<'_, C64>, y: &[C64], noise_var: f64) -> Result<Estimates> {
347 check(h, y, noise_var)?;
348 let (z, var) = self.single(h, y, noise_var)?;
349 if let Some(z) = z {
350 return Ok(Estimates { symbols: vec![z], noise_var: vec![var] });
351 }
352 let f = filter(h, self.criterion, self.constellation.energy(), noise_var)?;
353 let symbols = matvec(&f.w.view(), y).iter().zip(&f.bias).map(|(z, b)| z / b).collect();
354 Ok(Estimates { symbols, noise_var: f.var })
355 }
356
357 fn single(&self, h: &ArrayView2<'_, C64>, y: &[C64], noise_var: f64) -> Result<(Option<C64>, f64)> {
359 if h.ncols() != 1 {
360 return Ok((None, 0.0));
361 }
362 let es = self.constellation.energy();
363 let s = SingleStream::new(h, self.criterion, es, noise_var, regularisation(h, self.criterion, es, noise_var)?)?;
364 Ok((Some(s.estimate(h, y)), s.var))
365 }
366}
367
368impl MimoDetector for Linear {
369 fn constellation(&self) -> &Constellation {
370 &self.constellation
371 }
372
373 fn has_soft_output(&self) -> bool {
374 true
375 }
376
377 fn soft_estimate(&self, h: &ArrayView2<'_, C64>, y: &[C64], noise_var: f64) -> Result<Option<Estimates>> {
378 self.estimate(h, y, noise_var).map(Some)
379 }
380
381 fn detect(&self, h: &ArrayView2<'_, C64>, y: &[C64], noise_var: f64) -> Result<Vec<usize>> {
382 Ok(self.constellation.decide(&self.estimate(h, y, noise_var)?.symbols))
383 }
384
385 fn llr(&self, h: &ArrayView2<'_, C64>, y: &[C64], noise_var: f64) -> Result<Vec<f64>> {
386 let e = self.estimate(h, y, noise_var)?;
387 self.constellation.llr_with_variances(&e.symbols, &e.noise_var, None, Jacobian::Exact)
388 }
389
390 fn detect_llr(&self, h: &ArrayView2<'_, C64>, y: &[C64], noise_var: f64) -> Result<(Vec<usize>, Vec<f64>)> {
391 let e = self.estimate(h, y, noise_var)?;
392 let llr = self.constellation.llr_with_variances(&e.symbols, &e.noise_var, None, Jacobian::Exact)?;
393 Ok((self.constellation.decide(&e.symbols), llr))
394 }
395
396 fn detect_block(&self, h: &ArrayView2<'_, C64>, y: &ArrayView2<'_, C64>, noise_var: f64) -> Result<Array2<usize>> {
397 check_block(h, y, noise_var)?;
398 let f = filter(h, self.criterion, self.constellation.energy(), noise_var)?;
399 let z = f.w.dot(y);
400 Ok(Array2::from_shape_fn(z.dim(), |(k, t)| self.constellation.nearest(z[[k, t]] / f.bias[k])))
401 }
402}
403
404#[derive(Clone, Debug, PartialEq)]
407pub struct Sic {
408 pub constellation: Constellation,
410 pub criterion: Criterion,
412 pub ordered: bool,
414}
415
416impl Sic {
417 pub fn mmse(constellation: Constellation) -> Sic {
419 Sic { constellation, criterion: Criterion::Mmse, ordered: true }
420 }
421
422 pub fn zf(constellation: Constellation) -> Sic {
424 Sic { constellation, criterion: Criterion::Zf, ordered: true }
425 }
426}
427
428struct SicStage {
431 stream: usize,
432 w: Vec<C64>,
433 bias: f64,
434 column: Vec<C64>,
435}
436
437impl Sic {
438 fn stages(&self, h: &ArrayView2<'_, C64>, noise_var: f64) -> Result<Vec<SicStage>> {
440 let es = self.constellation.energy();
441 let mut remaining: Vec<usize> = (0..h.ncols()).collect();
442 let mut stages = Vec::with_capacity(h.ncols());
443 while !remaining.is_empty() {
444 let hc = Array2::from_shape_fn((h.nrows(), remaining.len()), |(i, j)| h[[i, remaining[j]]]);
445 let f = filter(&hc.view(), self.criterion, es, noise_var)?;
446 let k = if self.ordered {
447 (0..remaining.len()).min_by(|a, b| f.var[*a].total_cmp(&f.var[*b])).unwrap_or(0)
448 } else {
449 0
450 };
451 stages.push(SicStage {
452 stream: remaining[k],
453 w: f.w.row(k).to_vec(),
454 bias: f.bias[k],
455 column: hc.column(k).to_vec(),
456 });
457 remaining.remove(k);
458 }
459 Ok(stages)
460 }
461
462 fn cancel(&self, stages: &[SicStage], residual: &mut [C64]) -> Vec<usize> {
464 let mut out = vec![0; stages.len()];
465 for s in stages {
466 let z: C64 = s.w.iter().zip(residual.iter()).map(|(w, v)| w * v).sum::<C64>() / s.bias;
467 let idx = self.constellation.nearest(z);
468 let p = self.constellation.points()[idx];
469 for (r, hk) in residual.iter_mut().zip(&s.column) {
470 *r -= hk * p;
471 }
472 out[s.stream] = idx;
473 }
474 out
475 }
476}
477
478impl MimoDetector for Sic {
479 fn constellation(&self) -> &Constellation {
480 &self.constellation
481 }
482
483 fn detect(&self, h: &ArrayView2<'_, C64>, y: &[C64], noise_var: f64) -> Result<Vec<usize>> {
484 check(h, y, noise_var)?;
485 let stages = self.stages(h, noise_var)?;
486 Ok(self.cancel(&stages, &mut y.to_vec()))
487 }
488
489 fn detect_block(&self, h: &ArrayView2<'_, C64>, y: &ArrayView2<'_, C64>, noise_var: f64) -> Result<Array2<usize>> {
490 check_block(h, y, noise_var)?;
491 let stages = self.stages(h, noise_var)?;
492 by_column(h.ncols(), y, |col| Ok(self.cancel(&stages, col)))
493 }
494}
495
496#[derive(Clone, Debug, PartialEq)]
498pub struct Ml {
499 pub constellation: Constellation,
501}
502
503impl Ml {
504 pub fn new(constellation: Constellation) -> Ml {
506 Ml { constellation }
507 }
508
509 fn products(&self, h: &ArrayView2<'_, C64>) -> Result<Vec<C64>> {
516 let (m, nt, nr) = (self.constellation.order(), h.ncols(), h.nrows());
517 let total = (0..nt).try_fold(1usize, |acc, _| acc.checked_mul(m)).filter(|t| *t <= ML_MAX_HYPOTHESES);
518 ensure(total.is_some(), || format!("{m}^{nt} hypotheses exceed the limit of the exhaustive search"))?;
519 let pts = self.constellation.points();
520 let mut prod = Vec::with_capacity(nt * m * nr);
521 for k in 0..nt {
522 for p in pts {
523 prod.extend((0..nr).map(|r| h[[r, k]] * p));
524 }
525 }
526 Ok(prod)
527 }
528
529 fn search(&self, prod: &[C64], nt: usize, y: &[C64], mut visit: impl FnMut(&[usize], f64)) {
533 let (m, nr) = (self.constellation.order(), y.len());
534 let mut idx = vec![0usize; nt];
535 loop {
536 let d: f64 = (0..nr)
537 .map(|r| {
538 let a: C64 = idx.iter().enumerate().map(|(k, &i)| prod[(k * m + i) * nr + r]).sum();
539 (y[r] - a).norm_sqr()
540 })
541 .sum();
542 visit(&idx, d);
543 let mut k = 0;
545 loop {
546 if k == nt {
547 return;
548 }
549 idx[k] += 1;
550 if idx[k] < m {
551 break;
552 }
553 idx[k] = 0;
554 k += 1;
555 }
556 }
557 }
558
559 fn best(&self, prod: &[C64], nt: usize, y: &[C64]) -> Vec<usize> {
561 let mut best = (f64::INFINITY, vec![0; nt]);
562 self.search(prod, nt, y, |idx, d| {
563 if d < best.0 {
564 best.0 = d;
565 best.1.copy_from_slice(idx);
566 }
567 });
568 best.1
569 }
570}
571
572fn max_log(best: &[[f64; 2]], noise_var: f64, clip: f64) -> Vec<f64> {
574 best.iter()
575 .map(|[d0, d1]| match (d0.is_finite(), d1.is_finite()) {
576 (true, true) => ((d1 - d0) / noise_var).clamp(-clip, clip),
577 (true, false) => clip,
578 (false, true) => -clip,
579 (false, false) => 0.0,
580 })
581 .collect()
582}
583
584fn update_bits(best: &mut [[f64; 2]], idx: &[usize], bits: usize, d: f64) {
586 for (s, &i) in idx.iter().enumerate() {
587 for p in 0..bits {
588 let b = (i >> (bits - 1 - p)) & 1;
589 let slot = &mut best[s * bits + p][b];
590 if d < *slot {
591 *slot = d;
592 }
593 }
594 }
595}
596
597impl MimoDetector for Ml {
598 fn constellation(&self) -> &Constellation {
599 &self.constellation
600 }
601
602 fn has_soft_output(&self) -> bool {
603 true
604 }
605
606 fn detect(&self, h: &ArrayView2<'_, C64>, y: &[C64], noise_var: f64) -> Result<Vec<usize>> {
607 check(h, y, noise_var)?;
608 Ok(self.best(&self.products(h)?, h.ncols(), y))
609 }
610
611 fn llr(&self, h: &ArrayView2<'_, C64>, y: &[C64], noise_var: f64) -> Result<Vec<f64>> {
612 Ok(self.detect_llr(h, y, noise_var)?.1)
613 }
614
615 fn detect_llr(&self, h: &ArrayView2<'_, C64>, y: &[C64], noise_var: f64) -> Result<(Vec<usize>, Vec<f64>)> {
616 check(h, y, noise_var)?;
617 let prod = self.products(h)?;
618 let (nt, bits) = (h.ncols(), self.constellation.bits_per_symbol());
619 let mut best = (f64::INFINITY, vec![0; nt]);
620 let mut per_bit = vec![[f64::INFINITY; 2]; nt * bits];
621 self.search(&prod, nt, y, |idx, d| {
622 if d < best.0 {
623 best.0 = d;
624 best.1.copy_from_slice(idx);
625 }
626 update_bits(&mut per_bit, idx, bits, d);
627 });
628 Ok((best.1, max_log(&per_bit, noise_var, f64::INFINITY)))
629 }
630
631 fn detect_block(&self, h: &ArrayView2<'_, C64>, y: &ArrayView2<'_, C64>, noise_var: f64) -> Result<Array2<usize>> {
632 check_block(h, y, noise_var)?;
633 let prod = self.products(h)?;
634 by_column(h.ncols(), y, |col| Ok(self.best(&prod, h.ncols(), col)))
635 }
636}
637
638struct Triangular {
641 r: CMat,
642 qh: CMat,
643}
644
645impl Triangular {
646 fn new(h: &ArrayView2<'_, C64>, what: &str) -> Result<Triangular> {
648 tall(h, what)?;
649 let (q, r) = qr(h)?;
650 if (0..r.nrows()).any(|i| r[[i, i]].re <= 0.0) {
651 return Err(Error::Numerical("the channel is rank deficient".into()));
652 }
653 Ok(Triangular { r, qh: hermitian(&q.view()) })
654 }
655
656 fn rotate(&self, y: &[C64]) -> Vec<C64> {
658 matvec(&self.qh.view(), y)
659 }
660}
661
662type Frame = (Vec<(f64, usize)>, usize, f64);
665
666#[derive(Clone, Debug, PartialEq)]
672pub struct Sphere {
673 pub constellation: Constellation,
675}
676
677impl Sphere {
678 pub fn new(constellation: Constellation) -> Sphere {
680 Sphere { constellation }
681 }
682
683 fn search(&self, r: &CMat, z: &[C64]) -> Vec<usize> {
685 let n = r.ncols();
686 let pts = self.constellation.points();
687 let mut best = (f64::INFINITY, vec![0usize; n]);
688 let mut x = vec![C64::new(0.0, 0.0); n];
689 let mut idx = vec![0usize; n];
690 let mut frames: Vec<Frame> = Vec::with_capacity(n);
693 let mut spare: Vec<Vec<(f64, usize)>> = Vec::with_capacity(n);
694 let candidates = |k: usize, x: &[C64], mut c: Vec<(f64, usize)>| -> Vec<(f64, usize)> {
695 let mut b = z[k];
696 for j in k + 1..n {
697 b -= r[[k, j]] * x[j];
698 }
699 c.clear();
700 c.extend(pts.iter().enumerate().map(|(i, p)| ((b - r[[k, k]] * p).norm_sqr(), i)));
701 c.sort_by(|a, b| a.0.total_cmp(&b.0));
702 c
703 };
704 frames.push((candidates(n - 1, &x, Vec::with_capacity(pts.len())), 0, 0.0));
705 while !frames.is_empty() {
706 let k = n - frames.len();
707 let Some(f) = frames.last_mut() else { break };
708 let candidate = f.0.get(f.1).copied();
709 f.1 += 1;
710 let base = f.2;
711 let next = match candidate {
712 Some((d, i)) if base + d < best.0 => Some((base + d, i)),
713 _ => None,
715 };
716 let Some((partial, i)) = next else {
717 if let Some(f) = frames.pop() {
718 spare.push(f.0);
719 }
720 continue;
721 };
722 x[k] = pts[i];
723 idx[k] = i;
724 if k == 0 {
725 best.0 = partial;
726 best.1.copy_from_slice(&idx);
727 } else {
728 let c = candidates(k - 1, &x, spare.pop().unwrap_or_default());
729 frames.push((c, 0, partial));
730 }
731 }
732 best.1
733 }
734}
735
736impl MimoDetector for Sphere {
737 fn constellation(&self) -> &Constellation {
738 &self.constellation
739 }
740
741 fn detect(&self, h: &ArrayView2<'_, C64>, y: &[C64], noise_var: f64) -> Result<Vec<usize>> {
742 check(h, y, noise_var)?;
743 let t = Triangular::new(h, "the sphere decoder")?;
744 Ok(self.search(&t.r, &t.rotate(y)))
745 }
746
747 fn detect_block(&self, h: &ArrayView2<'_, C64>, y: &ArrayView2<'_, C64>, noise_var: f64) -> Result<Array2<usize>> {
748 check_block(h, y, noise_var)?;
749 let t = Triangular::new(h, "the sphere decoder")?;
750 by_column(h.ncols(), y, |col| Ok(self.search(&t.r, &t.rotate(col))))
751 }
752}
753
754#[derive(Clone, Debug, PartialEq)]
757pub struct KBest {
758 constellation: Constellation,
760 k: usize,
762 llr_clip: f64,
764}
765
766struct Paths {
769 labels: Vec<usize>,
770 dist: Vec<f64>,
771}
772
773impl Paths {
774 fn iter(&self, nt: usize) -> impl Iterator<Item = (&[usize], f64)> {
775 self.labels.chunks_exact(nt.max(1)).zip(self.dist.iter().copied())
777 }
778
779 fn first(&self, nt: usize) -> Vec<usize> {
780 self.labels.get(..nt).map(<[usize]>::to_vec).unwrap_or_default()
781 }
782}
783
784impl KBest {
785 pub fn new(constellation: Constellation, k: usize) -> Result<KBest> {
791 ensure(k >= 1, || "K-best needs k ≥ 1".into())?;
792 Ok(KBest { constellation, k, llr_clip: KNOWN_BIT_LLR })
793 }
794
795 pub fn with_llr_clip(mut self, llr_clip: f64) -> Result<KBest> {
801 ensure(llr_clip > 0.0, || format!("the LLR clip must be positive, got {llr_clip}"))?;
802 self.llr_clip = llr_clip;
803 Ok(self)
804 }
805
806 pub fn k(&self) -> usize {
808 self.k
809 }
810
811 pub fn llr_clip(&self) -> f64 {
813 self.llr_clip
814 }
815
816 fn search(&self, r: &CMat, z: &[C64]) -> Paths {
818 let n = r.ncols();
819 let pts = self.constellation.points();
820 let mut labels: Vec<usize> = Vec::new();
823 let mut next_labels: Vec<usize> = Vec::new();
824 let mut dist = vec![0.0];
825 let mut children: Vec<(f64, usize, usize)> = Vec::new();
828 for (depth, k) in (0..n).rev().enumerate() {
829 children.clear();
830 for (p, d) in dist.iter().enumerate() {
831 let mut b = z[k];
832 for (level, &l) in labels[p * depth..(p + 1) * depth].iter().enumerate() {
833 let j = n - 1 - level;
834 b -= r[[k, j]] * pts[l];
835 }
836 children.extend(pts.iter().enumerate().map(|(i, s)| (d + (b - r[[k, k]] * s).norm_sqr(), p, i)));
837 }
838 children.sort_by(|a, b| a.0.total_cmp(&b.0));
839 children.truncate(self.k);
840 next_labels.clear();
841 dist.clear();
842 for &(d, p, i) in &children {
843 next_labels.extend_from_slice(&labels[p * depth..(p + 1) * depth]);
844 next_labels.push(i);
845 dist.push(d);
846 }
847 std::mem::swap(&mut labels, &mut next_labels);
848 }
849 labels.chunks_exact_mut(n.max(1)).for_each(<[usize]>::reverse);
850 Paths { labels, dist }
851 }
852
853 fn soft(&self, paths: &Paths, nt: usize, noise_var: f64) -> Vec<f64> {
855 let bits = self.constellation.bits_per_symbol();
856 let mut best = vec![[f64::INFINITY; 2]; nt * bits];
857 for (labels, d) in paths.iter(nt) {
858 update_bits(&mut best, labels, bits, d);
859 }
860 max_log(&best, noise_var, self.llr_clip)
861 }
862}
863
864impl MimoDetector for KBest {
865 fn constellation(&self) -> &Constellation {
866 &self.constellation
867 }
868
869 fn has_soft_output(&self) -> bool {
870 true
871 }
872
873 fn detect(&self, h: &ArrayView2<'_, C64>, y: &[C64], noise_var: f64) -> Result<Vec<usize>> {
874 check(h, y, noise_var)?;
875 let t = Triangular::new(h, "K-best")?;
876 Ok(self.search(&t.r, &t.rotate(y)).first(h.ncols()))
877 }
878
879 fn llr(&self, h: &ArrayView2<'_, C64>, y: &[C64], noise_var: f64) -> Result<Vec<f64>> {
880 Ok(self.detect_llr(h, y, noise_var)?.1)
881 }
882
883 fn detect_llr(&self, h: &ArrayView2<'_, C64>, y: &[C64], noise_var: f64) -> Result<(Vec<usize>, Vec<f64>)> {
884 check(h, y, noise_var)?;
885 let t = Triangular::new(h, "K-best")?;
886 let paths = self.search(&t.r, &t.rotate(y));
887 Ok((paths.first(h.ncols()), self.soft(&paths, h.ncols(), noise_var)))
888 }
889
890 fn detect_block(&self, h: &ArrayView2<'_, C64>, y: &ArrayView2<'_, C64>, noise_var: f64) -> Result<Array2<usize>> {
891 check_block(h, y, noise_var)?;
892 let t = Triangular::new(h, "K-best")?;
893 by_column(h.ncols(), y, |col| Ok(self.search(&t.r, &t.rotate(col)).first(h.ncols())))
894 }
895}
896
897#[derive(Clone, Debug, PartialEq)]
901pub struct Coherent<D> {
902 pub detector: D,
904 pub h: CMat,
906}
907
908impl<D: MimoDetector> Coherent<D> {
909 pub fn new(detector: D, h: CMat) -> Coherent<D> {
911 Coherent { detector, h }
912 }
913}
914
915impl<D: MimoDetector> Detect for Coherent<D> {
916 fn detect_into(&mut self, y: &Symbols, noise_var: f64, out: &mut Symbols) -> Result<()> {
919 let labels = self.detector.detect_block(&self.h.view(), &y.as_array().view(), noise_var)?;
920 let pts = self.detector.constellation().points();
921 out.assign(&labels.mapv(|i| pts[i]));
922 Ok(())
923 }
924}
925
926#[cfg(test)]
927mod tests {
928 use super::*;
929 use comnum_core::rng::Rng;
930
931 fn channel(nr: usize, nt: usize, rng: &mut Rng) -> CMat {
932 Array2::from_shape_fn((nr, nt), |_| rng.complex_gaussian())
933 }
934
935 fn noisy(h: &CMat, x: &[C64], s2: f64, rng: &mut Rng) -> Vec<C64> {
936 matvec(&h.view(), x).into_iter().map(|v| v + rng.complex_gaussian() * s2.sqrt()).collect()
937 }
938
939 #[test]
940 fn the_matched_filter_leaves_the_other_streams_as_interference() {
941 let mut rng = Rng::new(11);
942 let (es, s2) = (2.0, 0.3);
943 let h = channel(6, 3, &mut rng);
944 let f = filter(&h.view(), Criterion::Mr, es, s2).unwrap();
945 let g = hermitian(&h.view()).dot(&h);
946 for k in 0..3 {
947 let wk = f.w.row(k);
950 let gain = |j: usize| wk.iter().zip(h.column(j)).map(|(a, b)| a * b).sum::<C64>();
951 assert!((gain(k) - 1.0).norm() < 1e-14 && f.bias[k] == 1.0);
952 let expected = s2 * wk.iter().map(C64::norm_sqr).sum::<f64>()
953 + es * (0..3).filter(|j| *j != k).map(|j| gain(j).norm_sqr()).sum::<f64>();
954 assert!((f.var[k] - expected).abs() < 1e-12 * expected);
955 assert!(
956 (f.var[k]
957 - (s2 * g[[k, k]].re + es * (g[[k, (k + 1) % 3]].norm_sqr() + g[[k, (k + 2) % 3]].norm_sqr()))
958 / g[[k, k]].re.powi(2))
959 .abs()
960 < 1e-12
961 );
962 }
963 let mut q = CMat::zeros((4, 2));
965 q[[0, 0]] = C64::new(2.0, 0.0);
966 q[[1, 1]] = C64::new(0.0, 3.0);
967 let (mr, zf) =
968 (filter(&q.view(), Criterion::Mr, es, s2).unwrap(), filter(&q.view(), Criterion::Zf, es, s2).unwrap());
969 assert!(mr.w.iter().zip(zf.w.iter()).all(|(a, b)| (a - b).norm() < 1e-15));
970 assert!(mr.var.iter().zip(&zf.var).all(|(a, b)| (a - b).abs() < 1e-15));
971 }
972
973 #[test]
974 fn the_single_stream_filter_is_the_matrix_filter() {
975 let mut rng = Rng::new(7);
976 let es = 1.0;
977 for nr in [1, 3] {
978 let h = channel(nr, 1, &mut rng);
979 for (criterion, lambda) in [(Criterion::Zf, 0.0), (Criterion::Mmse, 0.3 / es)] {
980 let f = filter(&h.view(), criterion, es, 0.3).unwrap();
981 let hh = hermitian(&h.view());
983 let ginv = inverse_hpd(&(hh.dot(&h) + identity(1).mapv(|x| x * lambda)).view()).unwrap();
984 let w = ginv.dot(&hh);
985 assert!(f.w.iter().zip(w.iter()).all(|(a, b)| (a - b).norm() < 1e-14 * b.norm().max(1.0)));
986 assert!(
987 (f.var[0]
988 - match criterion {
989 Criterion::Mr | Criterion::Zf => 0.3 * ginv[[0, 0]].re,
990 Criterion::Mmse => {
991 let e = 0.3 * ginv[[0, 0]].re;
992 es * e / (es - e)
993 }
994 })
995 .abs()
996 < 1e-14
997 );
998 }
999 }
1000 let qpsk = Constellation::psk(4).unwrap().normalized();
1002 let h = channel(3, 1, &mut rng);
1003 let y: Vec<C64> = (0..3).map(|_| rng.complex_gaussian()).collect();
1004 for d in [Linear::mr(qpsk.clone()), Linear::zf(qpsk.clone()), Linear::mmse(qpsk.clone())] {
1005 let f = filter(&h.view(), d.criterion, qpsk.energy(), 0.3).unwrap();
1006 let e = d.estimate(&h.view(), &y, 0.3).unwrap();
1007 assert!((e.symbols[0] - matvec(&f.w.view(), &y)[0] / f.bias[0]).norm() < 1e-14 && e.noise_var == f.var);
1008 }
1009 let zero = Array2::from_elem((2, 1), C64::new(0.0, 0.0));
1010 assert!(filter(&zero.view(), Criterion::Zf, es, 0.3).is_err());
1011 }
1012
1013 #[test]
1014 fn every_detector_is_exact_without_noise() {
1015 let c = Constellation::qam(16).unwrap().normalized();
1016 let mut rng = Rng::new(1);
1017 for _ in 0..50 {
1018 let h = channel(4, 3, &mut rng);
1019 let idx: Vec<usize> = (0..3).map(|_| rng.below(16)).collect();
1020 let x: Vec<C64> = idx.iter().map(|&i| c.points()[i]).collect();
1021 let y = matvec(&h.view(), &x);
1022 let dets: Vec<Box<dyn MimoDetector>> = vec![
1023 Box::new(Linear::zf(c.clone())),
1024 Box::new(Linear::mmse(c.clone())),
1025 Box::new(Sic::zf(c.clone())),
1026 Box::new(Sic::mmse(c.clone())),
1027 Box::new(Ml::new(c.clone())),
1028 Box::new(Sphere::new(c.clone())),
1029 Box::new(KBest::new(c.clone(), 4).unwrap()),
1030 ];
1031 for d in &dets {
1032 assert_eq!(d.detect(&h.view(), &y, 1e-6).unwrap(), idx);
1033 assert_eq!(d.constellation(), &c);
1034 }
1035 }
1036 }
1037
1038 #[test]
1039 fn the_sphere_decoder_is_the_exhaustive_ml_detector() {
1040 let c = Constellation::qam(4).unwrap().normalized();
1041 let mut rng = Rng::new(2);
1042 for _ in 0..300 {
1043 let h = channel(3, 3, &mut rng);
1044 let x: Vec<C64> = (0..3).map(|_| c.points()[rng.below(4)]).collect();
1045 let y = noisy(&h, &x, 0.5, &mut rng);
1046 let ml = Ml::new(c.clone()).detect(&h.view(), &y, 0.5).unwrap();
1047 assert_eq!(Sphere::new(c.clone()).detect(&h.view(), &y, 0.5).unwrap(), ml);
1048 assert_eq!(KBest::new(c.clone(), 64).unwrap().detect(&h.view(), &y, 0.5).unwrap(), ml);
1050 }
1051 }
1052
1053 #[test]
1054 fn soft_outputs_agree_in_sign_with_the_decisions() {
1055 let c = Constellation::qam(16).unwrap().normalized();
1056 let mut rng = Rng::new(3);
1057 let h = channel(4, 2, &mut rng);
1058 let x = [c.points()[5], c.points()[10]];
1059 let y = noisy(&h, &x, 1e-3, &mut rng);
1060 for d in
1061 [&Linear::mmse(c.clone()) as &dyn MimoDetector, &Ml::new(c.clone()), &KBest::new(c.clone(), 8).unwrap()]
1062 {
1063 assert!(d.has_soft_output());
1064 let l = d.llr(&h.view(), &y, 1e-3).unwrap();
1065 assert_eq!(l.len(), 8);
1066 let bits: Vec<usize> = [5usize, 10].iter().flat_map(|i| (0..4).map(move |p| (i >> (3 - p)) & 1)).collect();
1067 for (v, b) in l.iter().zip(bits) {
1068 assert!(if b == 0 { *v > 0.0 } else { *v < 0.0 }, "{l:?}");
1069 }
1070 }
1071 assert!(Sphere::new(c.clone()).llr(&h.view(), &y, 1e-3).is_err());
1072 assert!(!Sphere::new(c.clone()).has_soft_output() && !Sic::mmse(c.clone()).has_soft_output());
1073 assert!(KBest::new(c.clone(), 0).is_err() && KBest::new(c, 2).unwrap().with_llr_clip(-1.0).is_err());
1074 }
1075
1076 #[test]
1077 fn linear_block_detection_matches_vector_detection() {
1078 let c = Constellation::qam(4).unwrap().normalized();
1079 let mut rng = Rng::new(4);
1080 let h = channel(2, 2, &mut rng);
1081 let y = Array2::from_shape_fn((2, 20), |_| rng.complex_gaussian());
1082 for d in [Linear::zf(c.clone()), Linear::mmse(c.clone())] {
1083 let block = d.detect_block(&h.view(), &y.view(), 0.1).unwrap();
1084 for t in 0..20 {
1085 assert_eq!(block.column(t).to_vec(), d.detect(&h.view(), &y.column(t).to_vec(), 0.1).unwrap());
1086 }
1087 }
1088 }
1089
1090 #[test]
1093 fn block_and_joint_outputs_match_the_detection_of_every_use() {
1094 let c = Constellation::qam(16).unwrap().normalized();
1095 let mut rng = Rng::new(8);
1096 let dets: Vec<Box<dyn MimoDetector>> = vec![
1097 Box::new(Sic::zf(c.clone())),
1098 Box::new(Sic::mmse(c.clone())),
1099 Box::new(Sic { constellation: c.clone(), criterion: Criterion::Mmse, ordered: false }),
1100 Box::new(Ml::new(c.clone())),
1101 Box::new(Sphere::new(c.clone())),
1102 Box::new(KBest::new(c.clone(), 1).unwrap()),
1103 Box::new(KBest::new(c.clone(), 12).unwrap()),
1104 ];
1105 for _ in 0..10 {
1106 let h = channel(4, 3, &mut rng);
1107 let y = Array2::from_shape_fn((4, 8), |_| rng.complex_gaussian());
1108 for d in &dets {
1109 let block = d.detect_block(&h.view(), &y.view(), 0.2).unwrap();
1110 for t in 0..8 {
1111 let col = y.column(t).to_vec();
1112 let labels = d.detect(&h.view(), &col, 0.2).unwrap();
1113 assert_eq!(block.column(t).to_vec(), labels);
1114 if let Ok(llr) = d.llr(&h.view(), &col, 0.2) {
1115 let (l, v) = d.detect_llr(&h.view(), &col, 0.2).unwrap();
1116 assert_eq!(
1117 (l, v.iter().map(|x| x.to_bits()).collect::<Vec<_>>()),
1118 (labels, llr.iter().map(|x| x.to_bits()).collect())
1119 );
1120 }
1121 }
1122 }
1123 }
1124 }
1125
1126 #[test]
1127 fn invalid_inputs_are_errors() {
1128 let c = Constellation::qam(4).unwrap().normalized();
1129 let h = channel(2, 3, &mut Rng::new(5));
1130 let y = vec![C64::new(0.0, 0.0); 2];
1131 assert!(Linear::zf(c.clone()).detect(&h.view(), &y, 0.1).is_err());
1132 assert!(Sphere::new(c.clone()).detect(&h.view(), &y, 0.1).is_err());
1133 assert!(Linear::mmse(c.clone()).detect(&h.view(), &y, 0.1).is_ok()); assert!(Ml::new(c.clone()).detect(&h.view(), &y[..1], 0.1).is_err());
1135 assert!(Ml::new(c).detect(&h.view(), &y, 0.0).is_err());
1136 }
1137
1138 #[test]
1139 fn the_coherent_block_implements_detect() {
1140 let c = Constellation::qam(4).unwrap().normalized();
1141 let mut rng = Rng::new(6);
1142 let h = channel(3, 2, &mut rng);
1143 let x = Array2::from_shape_fn((2, 8), |_| c.points()[rng.below(4)]);
1144 let y = Symbols::new(h.dot(&x));
1145 let mut d = Coherent::new(Sphere::new(c), h);
1146 assert_eq!(d.detect_into_owned(&y).unwrap().as_array(), &x);
1147 }
1148
1149 trait Owned {
1150 fn detect_into_owned(&mut self, y: &Symbols) -> Result<Symbols>;
1151 }
1152
1153 impl<D: MimoDetector> Owned for Coherent<D> {
1154 fn detect_into_owned(&mut self, y: &Symbols) -> Result<Symbols> {
1155 let mut out = Symbols::default();
1156 self.detect_into(y, 1e-3, &mut out)?;
1157 Ok(out)
1158 }
1159 }
1160}