1use crate::check_rate;
61use comnum_core::channel::Path;
62use comnum_core::prelude::*;
63use comnum_core::{ensure, ensure_shape};
64use comnum_dsp::signal::{extend_with_prefix, strip_prefix};
65use comnum_math::linalg::CMat;
66use comnum_math::{C64, cis, fft};
67
68#[derive(Clone, Copy, Debug, PartialEq)]
71pub struct Otfs {
72 pub n_fft: usize,
74 pub n_symbols: usize,
76 pub n_cp: usize,
79 pub sample_rate: Hertz,
82}
83
84impl Otfs {
85 pub fn new(n_fft: usize, n_symbols: usize) -> Result<Otfs> {
90 let o = Otfs { n_fft, n_symbols, n_cp: 0, sample_rate: Hertz(1.0) };
91 o.check()?;
92 Ok(o)
93 }
94
95 pub fn check(&self) -> Result<()> {
101 ensure(self.n_fft >= 1 && self.n_symbols >= 1, || {
102 format!("an OTFS grid needs at least one bin, got {} × {}", self.n_fft, self.n_symbols)
103 })?;
104 ensure(self.n_cp <= self.n_fft * self.n_symbols, || {
105 format!("a prefix of {} samples is longer than the frame of {}", self.n_cp, self.n_fft * self.n_symbols)
106 })?;
107 check_rate(self.sample_rate)
108 }
109
110 pub fn block_len(&self) -> usize {
112 self.n_fft * self.n_symbols + self.n_cp
113 }
114
115 fn doppler_dft(&self, v: &[C64], inverse: bool) -> Vec<C64> {
117 let (m, n) = (self.n_fft, self.n_symbols);
118 let mut out = vec![C64::new(0.0, 0.0); m * n];
119 let mut col = vec![C64::new(0.0, 0.0); n];
120 for d in 0..m {
121 for (k, c) in col.iter_mut().enumerate() {
122 *c = v[k * m + d];
123 }
124 if inverse {
125 fft::ifft_unitary_in_place(&mut col);
126 } else {
127 fft::fft_unitary_in_place(&mut col);
128 }
129 for (k, c) in col.iter().enumerate() {
130 out[k * m + d] = *c;
131 }
132 }
133 out
134 }
135
136 pub fn modulate(&self, x: &[C64]) -> Result<Vec<C64>> {
143 self.check()?;
144 let mn = self.n_fft * self.n_symbols;
145 ensure_shape(x.len() % mn == 0, || format!("{} symbols for frames of {mn}", x.len()))?;
146 let mut out = Vec::with_capacity(x.len() / mn * self.block_len());
147 for frame in x.chunks_exact(mn) {
148 extend_with_prefix(&mut out, &self.doppler_dft(frame, true), self.n_cp)?;
149 }
150 Ok(out)
151 }
152
153 pub fn demodulate(&self, r: &[C64]) -> Result<Vec<C64>> {
159 self.check()?;
160 let len = self.block_len();
161 ensure_shape(r.len() % len == 0, || format!("{} samples for frames of {len}", r.len()))?;
162 let mut out = Vec::with_capacity(r.len() / len * (len - self.n_cp));
163 for frame in r.chunks_exact(len) {
164 out.extend(self.doppler_dft(strip_prefix(frame, self.n_cp)?, false));
165 }
166 Ok(out)
167 }
168
169 pub fn effective_channel(&self, paths: &[Path]) -> Result<EffectiveChannel> {
209 self.check()?;
210 let (m, n) = (self.n_fft, self.n_symbols);
211 let mn = m * n;
212 let n_paths = paths.len();
213 let mut cols = vec![0usize; mn * n_paths];
214 let mut vals = vec![C64::new(0.0, 0.0); mn * n_paths];
215 for (i, path) in paths.iter().enumerate() {
216 let l = path.delay;
217 ensure(l < m, || format!("a delay of {l} bins needs more than {m} delay bins"))?;
218 let bins = path.doppler_hz * mn as f64 / self.sample_rate.0;
219 ensure((bins - bins.round()).abs() < 1e-9, || {
220 format!("a Doppler shift of {bins} bins is not a whole bin: use transform_channel")
221 })?;
222 let k = bins.round() as i64;
223 let gain = path.gain * cis(std::f64::consts::TAU * k as f64 * self.n_cp as f64 / mn as f64);
224 for p in 0..mn {
225 let (row_n, row_m) = (p / m, p % m);
226 let shifted = (row_m + m - l) % m;
227 let q = shifted + m * ((row_n as i64 - k).rem_euclid(n as i64) as usize);
228 let z = cis(std::f64::consts::TAU * k as f64 * shifted as f64 / mn as f64);
229 let wrap =
230 if row_m < l { cis(-std::f64::consts::TAU * row_n as f64 / n as f64) } else { C64::new(1.0, 0.0) };
231 cols[p * n_paths + i] = q;
232 vals[p * n_paths + i] = gain * z * wrap;
233 }
234 }
235 Ok(EffectiveChannel { size: mn, n_paths, cols, vals })
236 }
237
238 pub fn transform_channel(&self, h: &CMat) -> Result<CMat> {
249 let mn = self.n_fft * self.n_symbols;
250 ensure_shape(h.nrows() == mn && h.ncols() == mn, || {
251 format!("a {}×{} channel for frames of {mn}", h.nrows(), h.ncols())
252 })?;
253 let mut right = CMat::zeros((mn, mn));
255 for (p, row) in h.rows().into_iter().enumerate() {
256 let t = self.doppler_dft(&row.to_vec(), true);
257 right.row_mut(p).iter_mut().zip(t).for_each(|(o, v)| *o = v);
258 }
259 let mut out = CMat::zeros((mn, mn));
261 for (q, col) in right.columns().into_iter().enumerate() {
262 let t = self.doppler_dft(&col.to_vec(), false);
263 out.column_mut(q).iter_mut().zip(t).for_each(|(o, v)| *o = v);
264 }
265 Ok(out)
266 }
267}
268
269#[derive(Clone, Debug, PartialEq)]
273pub struct EffectiveChannel {
274 size: usize,
275 n_paths: usize,
276 cols: Vec<usize>,
278 vals: Vec<C64>,
280}
281
282impl EffectiveChannel {
283 pub fn size(&self) -> usize {
285 self.size
286 }
287
288 pub fn n_paths(&self) -> usize {
290 self.n_paths
291 }
292
293 pub fn row(&self, p: usize) -> impl Iterator<Item = (usize, C64)> + '_ {
298 let r = p * self.n_paths..(p + 1) * self.n_paths;
299 self.cols[r.clone()].iter().copied().zip(self.vals[r].iter().copied())
300 }
301
302 pub fn apply(&self, x: &[C64]) -> Result<Vec<C64>> {
307 ensure_shape(x.len() == self.size, || format!("{} values for an effective channel of {}", x.len(), self.size))?;
308 Ok((0..self.size).map(|p| self.row(p).map(|(q, v)| v * x[q]).sum()).collect())
309 }
310
311 pub fn to_dense(&self) -> CMat {
314 let mut h = CMat::zeros((self.size, self.size));
315 for i in 0..self.n_paths {
317 for p in 0..self.size {
318 h[[p, self.cols[p * self.n_paths + i]]] += self.vals[p * self.n_paths + i];
319 }
320 }
321 h
322 }
323}
324
325impl Modulate for Otfs {
328 fn modulate_into(&self, x: &Symbols, out: &mut Signal) -> Result<()> {
329 out.set_stream(&Otfs::modulate(self, crate::one_stream(x, "OTFS")?)?, self.sample_rate);
330 Ok(())
331 }
332}
333
334impl Demodulate for Otfs {
336 fn demodulate_into(&mut self, y: &Signal, out: &mut Symbols) -> Result<()> {
340 crate::check_input(y, self.sample_rate, "OTFS")?;
341 out.set_stream(&Otfs::demodulate(self, y.stream(0))?);
342 Ok(())
343 }
344}
345
346#[cfg(test)]
347mod tests {
348 use super::*;
349
350 #[test]
351 fn the_traits_carry_the_rate_of_the_modem() {
352 let mut o = Otfs { n_cp: 2, sample_rate: Hertz(480e3), ..Otfs::new(8, 4).unwrap() };
353 let x = Symbols::from_stream((0..32).map(|i| C64::new(f64::from(i), 1.0)).collect()).unwrap();
354 let s = Modulate::modulate(&o, &x).unwrap();
355 assert_eq!(s.sample_rate(), Hertz(480e3));
356 assert!(
357 Demodulate::demodulate(&mut o, &s)
358 .unwrap()
359 .as_slice()
360 .iter()
361 .zip(x.as_slice())
362 .all(|(a, b)| (a - b).norm() < 1e-12)
363 );
364 let wrong = Signal::from_stream(s.stream(0).to_vec(), Hertz(1.0)).unwrap();
365 assert!(Demodulate::demodulate(&mut o, &wrong).is_err());
366 assert!(Otfs { sample_rate: Hertz(f64::NAN), ..o }.check().is_err());
367 }
368
369 #[test]
370 fn a_channel_without_delay_or_doppler_is_the_identity_and_rows_are_sparse() {
371 let o = Otfs { sample_rate: Hertz(24.0), ..Otfs::new(6, 4).unwrap() };
373 let h = o.effective_channel(&paths_of(&[(C64::new(0.5, 0.0), 0, 0.0)])).unwrap().to_dense();
374 for p in 0..24 {
375 for q in 0..24 {
376 assert_eq!(h[[p, q]], if p == q { C64::new(0.5, 0.0) } else { C64::new(0.0, 0.0) });
377 }
378 }
379 let paths = [(C64::new(1.0, 0.0), 0, 0.0), (C64::new(0.3, 0.2), 2, 1.0), (C64::new(-0.4, 0.1), 5, -2.0)];
380 let sparse = o.effective_channel(&paths_of(&paths)).unwrap();
381 let h = sparse.to_dense();
382 for p in 0..24 {
383 assert_eq!((0..24).filter(|&q| h[[p, q]].norm() > 0.0).count(), 3);
384 assert_eq!((0..24).filter(|&q| h[[q, p]].norm() > 0.0).count(), 3);
385 assert_eq!(sparse.row(p).count(), 3);
386 }
387 assert!(o.effective_channel(&paths_of(&[(C64::new(1.0, 0.0), 0, 0.5)])).is_err());
388 assert!(sparse.apply(&[C64::new(0.0, 0.0); 23]).is_err());
389 assert!(Otfs { n_cp: 25, ..o }.check().is_err());
390 }
391
392 fn paths_of(p: &[(C64, usize, f64)]) -> Vec<Path> {
394 p.iter().map(|&(gain, delay, doppler_hz)| Path { gain, delay, doppler_hz }).collect()
395 }
396
397 fn dense_theorem_1(o: &Otfs, paths: &[(C64, usize, f64)], sample_rate: f64) -> CMat {
400 let (m, n) = (o.n_fft, o.n_symbols);
401 let mn = m * n;
402 let mut h = CMat::zeros((mn, mn));
403 for &(gain, l, doppler_hz) in paths {
404 let k = (doppler_hz * mn as f64 / sample_rate).round() as i64;
405 let gain = gain * cis(std::f64::consts::TAU * k as f64 * o.n_cp as f64 / mn as f64);
406 for p in 0..mn {
407 let (row_n, row_m) = (p / m, p % m);
408 let shifted = (row_m + m - l) % m;
409 let q = shifted + m * ((row_n as i64 - k).rem_euclid(n as i64) as usize);
410 let z = cis(std::f64::consts::TAU * k as f64 * shifted as f64 / mn as f64);
411 let wrap =
412 if row_m < l { cis(-std::f64::consts::TAU * row_n as f64 / n as f64) } else { C64::new(1.0, 0.0) };
413 h[[p, q]] += gain * z * wrap;
414 }
415 }
416 h
417 }
418
419 #[test]
420 fn the_sparse_channel_is_the_former_dense_matrix() {
421 let mut rng = comnum_core::rng::Rng::new(8);
422 for (m, n, cp) in [(6usize, 4usize, 0usize), (16, 8, 4), (5, 7, 3)] {
423 let mn = (m * n) as f64;
424 let o = Otfs { n_cp: cp, sample_rate: Hertz(mn * 1e3), ..Otfs::new(m, n).unwrap() };
425 let paths = [
427 (rng.complex_gaussian(), 0usize, 0.0),
428 (rng.complex_gaussian(), 2, 1e3),
429 (rng.complex_gaussian(), m - 1, -2e3),
430 (rng.complex_gaussian(), 2, 1e3),
431 ];
432 let sparse = o.effective_channel(&paths_of(&paths)).unwrap();
433 assert_eq!(sparse.to_dense(), dense_theorem_1(&o, &paths, mn * 1e3), "{m} × {n}");
434 let x: Vec<C64> = (0..m * n).map(|_| rng.complex_gaussian()).collect();
435 let dense = sparse.to_dense();
436 for (p, y) in sparse.apply(&x).unwrap().iter().enumerate() {
437 let want: C64 = (0..m * n).map(|q| dense[[p, q]] * x[q]).sum();
438 assert!((y - want).norm() < 1e-12);
439 }
440 }
441 }
442
443 fn transform_channel_by_columns(o: &Otfs, h: &CMat) -> CMat {
446 let mn = o.n_fft * o.n_symbols;
447 let mut out = CMat::zeros((mn, mn));
448 for q in 0..mn {
449 let mut e = vec![C64::new(0.0, 0.0); mn];
450 e[q] = C64::new(1.0, 0.0);
451 let s = o.doppler_dft(&e, true);
452 let hs: Vec<C64> = (0..mn).map(|p| (0..mn).map(|j| h[[p, j]] * s[j]).sum()).collect();
453 for (p, v) in o.doppler_dft(&hs, false).into_iter().enumerate() {
454 out[[p, q]] = v;
455 }
456 }
457 out
458 }
459
460 #[test]
461 fn the_transformed_channel_is_the_former_product_and_theorem_1_for_whole_bins() {
462 let mut rng = comnum_core::rng::Rng::new(5);
463 for (m, n) in [(4usize, 3usize), (8, 4), (5, 7)] {
464 let o = Otfs::new(m, n).unwrap();
465 let mn = m * n;
466 let h = CMat::from_shape_fn((mn, mn), |_| rng.complex_gaussian());
467 let (fast, slow) = (o.transform_channel(&h).unwrap(), transform_channel_by_columns(&o, &h));
468 let err = fast.iter().zip(&slow).map(|(a, b)| (a - b).norm()).fold(0.0, f64::max);
469 assert!(err < 1e-12, "{m} × {n}: {err}");
470 }
471 let o = Otfs::new(6, 4).unwrap();
473 let paths = [(C64::new(1.0, 0.0), 0usize, 0.0), (C64::new(0.3, 0.2), 2, 1.0), (C64::new(-0.4, 0.1), 5, -2.0)];
474 let mut h = CMat::zeros((24, 24));
475 for &(g, l, k) in &paths {
476 for p in 0..24 {
477 h[[p, (p + 24 - l) % 24]] += g * cis(std::f64::consts::TAU * k * (p as f64 - l as f64) / 24.0);
478 }
479 }
480 let theorem = Otfs { sample_rate: Hertz(24.0), ..o };
481 let (a, b) =
482 (o.transform_channel(&h).unwrap(), theorem.effective_channel(&paths_of(&paths)).unwrap().to_dense());
483 assert!(a.iter().zip(&b).all(|(x, y)| (x - y).norm() < 1e-12));
484 }
485}
486
487#[cfg(test)]
488mod literature {
489 use super::*;
490
491 type Matrix = [[C64; 4]; 4];
492
493 #[test]
497 fn raviteja_2019_example_of_theorem_1() {
498 let o = Otfs { sample_rate: Hertz(4.0), ..Otfs::new(2, 2).unwrap() };
500 let (z0, one) = (C64::new(0.0, 0.0), C64::new(1.0, 0.0));
501 let (j4, mj4, mj2) =
502 (cis(std::f64::consts::FRAC_PI_2), cis(-std::f64::consts::FRAC_PI_2), cis(-std::f64::consts::PI));
503 let cases: [((usize, f64), Matrix); 4] = [
504 ((0, 0.0), [[one, z0, z0, z0], [z0, one, z0, z0], [z0, z0, one, z0], [z0, z0, z0, one]]),
505 ((1, 0.0), [[z0, one, z0, z0], [one, z0, z0, z0], [z0, z0, z0, mj2], [z0, z0, one, z0]]),
506 ((0, 1.0), [[z0, z0, one, z0], [z0, z0, z0, j4], [one, z0, z0, z0], [z0, j4, z0, z0]]),
507 ((1, 1.0), [[z0, z0, z0, j4], [z0, z0, one, z0], [z0, mj4, z0, z0], [one, z0, z0, z0]]),
508 ];
509 for ((l, k), want) in cases {
510 let t = o.effective_channel(&[Path { gain: one, delay: l, doppler_hz: k }]).unwrap().to_dense();
511 for p in 0..4 {
512 for q in 0..4 {
513 assert!((t[[p, q]] - want[p][q]).norm() < 1e-12, "l {l} k {k} ({p}, {q}): {}", t[[p, q]]);
514 }
515 }
516 }
517 }
518}