sachy.dev
Repo for designs & driver for a TA7642 powered lightning detector
1//! A generalised Embassy driver for the Strike Sensor v1
2//!
3#![no_std]
4
5#[cfg(feature = "alloc")]
6extern crate alloc;
7
8mod analysis;
9pub mod drivers;
10pub mod traits;
11
12use core::cell::Cell;
13
14#[cfg(feature = "debug")]
15use defmt::info;
16use embassy_sync::{
17 blocking_mutex::raw::NoopRawMutex,
18 zerocopy_channel::{Channel, Receiver, Sender},
19};
20use embassy_time::{Duration, Instant, Ticker, Timer};
21
22use crate::traits::{AdcSource, BufferMut, PwmSource, TimeSource};
23
24pub const BLOCK_SIZE: usize = 512;
25pub type ZeroCopyChannel<'device> = Channel<'device, NoopRawMutex, (i64, [u16; BLOCK_SIZE])>;
26
27#[derive(Debug)]
28#[cfg_attr(feature = "defmt", derive(defmt::Format))]
29struct DetectorState {
30 max_duty: Cell<u8>,
31 duty: Cell<u8>,
32 avg: Cell<u16>,
33 strikes: Cell<u16>,
34 warn_level: Cell<u16>,
35}
36
37impl Default for DetectorState {
38 fn default() -> Self {
39 Self {
40 max_duty: Cell::new(0),
41 duty: Cell::new(0),
42 avg: Cell::new(0),
43 strikes: Cell::new(0),
44 warn_level: Cell::new(255),
45 }
46 }
47}
48
49#[derive(Debug, Clone)]
50#[cfg_attr(feature = "defmt", derive(defmt::Format))]
51pub struct DetectorConfig {
52 blip_threshold: Cell<u16>,
53 blip_size: Cell<usize>,
54}
55
56impl DetectorConfig {
57 pub const fn new(blip_threshold: u16, blip_size: usize) -> Self {
58 Self {
59 blip_threshold: Cell::new(blip_threshold),
60 blip_size: Cell::new(blip_size),
61 }
62 }
63
64 pub const fn blip_threshold(&self) -> u16 {
65 self.blip_threshold.get()
66 }
67
68 pub const fn blip_size(&self) -> usize {
69 self.blip_size.get()
70 }
71
72 pub fn set_blip_threshold(&self, blip_threshold: u16) {
73 self.blip_threshold.set(blip_threshold);
74 }
75
76 pub fn set_blip_size(&self, blip_size: usize) {
77 self.blip_size.set(blip_size);
78 }
79}
80
81impl Default for DetectorConfig {
82 fn default() -> Self {
83 Self::new(14, 2)
84 }
85}
86
87#[derive(Debug, Clone, Copy, PartialEq, Eq)]
88#[cfg_attr(feature = "defmt", derive(defmt::Format))]
89pub enum DetectorUpdate<'a> {
90 Tick {
91 timestamp: i64,
92 level: u16,
93 },
94 Detection {
95 timestamp: i64,
96 average: u16,
97 samples: &'a [u16],
98 peaks: &'a [(usize, u16)],
99 },
100}
101
102pub struct DetectorDriver<T: TimeSource, P: PwmSource, A: AdcSource> {
103 state: DetectorState,
104 config: DetectorConfig,
105 timer: T,
106 pwm: P,
107 adc: A,
108}
109
110impl<T, P, A> DetectorDriver<T, P, A>
111where
112 T: TimeSource,
113 P: PwmSource,
114 A: AdcSource,
115{
116 pub fn new(config: DetectorConfig, timer: T, pwm: P, adc: A) -> Self {
117 Self {
118 state: Default::default(),
119 config,
120 timer,
121 pwm,
122 adc,
123 }
124 }
125
126 pub async fn reset_timer_source(&mut self, timer_source: T::Source) {
127 self.timer.reset_from_source(timer_source).await;
128 }
129
130 pub fn get_timestamp(&self) -> i64 {
131 self.timer.timestamp()
132 }
133
134 pub fn set_blip_threshold(&self, threshold: u16) {
135 self.config.blip_threshold.set(threshold);
136 }
137
138 pub fn set_blip_size(&self, size: usize) {
139 self.config.blip_size.set(size);
140 }
141
142 pub fn set_config(&self, config: DetectorConfig) {
143 self.set_blip_size(config.blip_size());
144 self.set_blip_threshold(config.blip_threshold());
145 }
146
147 pub fn get_config(&self) -> DetectorConfig {
148 self.config.clone()
149 }
150
151 pub async fn tune(&mut self, samples: &mut [u16]) {
152 let mut duty = 0;
153 let max_duty = self.pwm.max_duty();
154 self.pwm.set_duty(duty);
155 Timer::after_secs(2).await;
156 let mut act_value = self.adc.sample_average(samples).await;
157 #[cfg(feature = "debug")]
158 info!("START: {}", act_value);
159
160 while act_value < 1500 {
161 duty += 2;
162 if duty >= max_duty {
163 #[cfg(feature = "debug")]
164 info!("RESET: {} @ Duty {}", act_value, duty);
165 duty = 0;
166 self.pwm.set_duty(duty);
167 Timer::after_secs(2).await;
168 act_value = self.adc.sample_average(samples).await;
169 continue;
170 }
171 self.pwm.set_duty(duty);
172 Timer::after_millis(250).await;
173 act_value = self.adc.sample_average(samples).await;
174 #[cfg(feature = "debug")]
175 info!("UPDATE: {} @ Duty {}", act_value, duty);
176 }
177 self.state.max_duty.set(duty as u8);
178 duty = (duty / 6) * 5;
179 self.pwm.set_duty(duty);
180 self.state.duty.set(duty as u8);
181 // Allow voltage level to stabilize after tuning
182 Timer::after_secs(2).await;
183 let avg = self.adc.sample_average(samples).await;
184 #[cfg(feature = "debug")]
185 info!("SET: {} @ Duty {}/{}", avg, duty, self.state.max_duty.get());
186 self.state.avg.set(avg);
187 }
188
189 /// Samples and returns an `i64` timestamp.
190 pub async fn sample_with_zerocopy<'a>(
191 &self,
192 dma: &mut Sender<'a, NoopRawMutex, (i64, [u16; BLOCK_SIZE])>,
193 ) {
194 loop {
195 let (time, samples) = dma.send().await;
196 self.adc.sample(samples).await;
197 *time = self.get_timestamp();
198 dma.send_done();
199 }
200 }
201
202 pub fn detect_from_sample<B, F>(
203 &self,
204 timestamp: i64,
205 samples: &[u16],
206 peaks: &mut B,
207 update: F,
208 ) where
209 B: BufferMut<(usize, u16)>,
210 F: Fn(DetectorUpdate<'_>),
211 {
212 peaks.clear();
213
214 let new_avg = analysis::analyse_buffer_by_stepped_windows(
215 self.config.blip_threshold.get(),
216 samples,
217 self.state.avg.get(),
218 peaks,
219 );
220
221 let blips = peaks.len();
222
223 if blips >= self.config.blip_size.get() {
224 self.state
225 .strikes
226 .update(|strike| strike.saturating_add(32));
227
228 update(DetectorUpdate::Detection {
229 timestamp,
230 average: self.state.avg.get(),
231 samples,
232 peaks: peaks.as_slice(),
233 });
234 }
235
236 self.state.avg.set(new_avg);
237 }
238
239 pub async fn detect_with_zerocopy<'a, 'd, B, F>(
240 &self,
241 dma: &mut Receiver<'a, NoopRawMutex, (i64, [u16; BLOCK_SIZE])>,
242 peaks: &'d mut B,
243 update: F,
244 ) where
245 B: BufferMut<(usize, u16)>,
246 F: Fn(DetectorUpdate<'_>),
247 {
248 loop {
249 let (timestamp, samples) = dma.receive().await;
250
251 self.detect_from_sample(*timestamp, samples, peaks, &update);
252
253 dma.receive_done();
254 }
255 }
256
257 pub async fn tick<F>(&self, update: F)
258 where
259 F: Fn(DetectorUpdate),
260 {
261 let mut inactive: Option<Instant> = None;
262 let mut interval = Ticker::every(Duration::from_secs(1));
263
264 loop {
265 interval.next().await;
266 let strikes = self.state.strikes.get();
267 let mut warn_level = self.state.warn_level.get();
268
269 if strikes > 32 {
270 warn_level = warn_level.saturating_add(strikes);
271 }
272
273 let decay = warn_level >> 8;
274
275 self.state.strikes.set(0);
276 self.state.warn_level.set(warn_level - decay);
277
278 match inactive {
279 Some(_) if decay > 0 => {
280 inactive = None;
281 }
282 Some(val) if val.elapsed() >= Duration::from_secs(3600) => break,
283 None if decay == 0 => {
284 inactive = Some(Instant::now());
285 }
286 None => {
287 update(DetectorUpdate::Tick {
288 timestamp: self.get_timestamp(),
289 level: warn_level - 255,
290 });
291 }
292 _ => continue,
293 }
294 }
295 }
296}
297
298#[cfg(test)]
299mod tests {
300 use core::{future::poll_fn, ops::RangeBounds, slice::SliceIndex};
301 use embassy_time::MockDriver;
302 use rand::{Rng, distr::uniform::SampleRange};
303
304 #[cfg(not(feature = "alloc"))]
305 extern crate alloc;
306 use super::*;
307
308 #[derive(Debug, Default)]
309 struct MockMachine {
310 pwm: Cell<u16>,
311 }
312
313 impl MockMachine {
314 fn adc_sample_avg(&self) -> u16 {
315 self.pwm.get().saturating_mul(14)
316 }
317 }
318
319 struct MockPwm<'a>(&'a MockMachine);
320
321 impl PwmSource for MockPwm<'_> {
322 fn max_duty(&self) -> u16 {
323 255
324 }
325
326 fn set_duty(&mut self, duty: u16) {
327 self.0.pwm.set(duty);
328 }
329
330 // These do nothing
331 fn enable(&mut self) {}
332 fn disable(&mut self) {}
333 }
334
335 struct MockAdc<'a>(&'a MockMachine);
336
337 impl AdcSource for MockAdc<'_> {
338 async fn sample(&self, samples: &mut [u16]) {
339 samples.fill(self.0.adc_sample_avg());
340 }
341
342 async fn sample_average(&self, _samples: &mut [u16]) -> u16 {
343 self.0.adc_sample_avg()
344 }
345 }
346
347 struct MockTimeSource;
348
349 impl TimeSource for MockTimeSource {
350 type Source = ();
351
352 async fn reset_from_source(&mut self, _source: Self::Source) {}
353
354 fn timestamp(&self) -> i64 {
355 0
356 }
357 }
358
359 #[cfg(not(feature = "alloc"))]
360 impl BufferMut<(usize, u16)> for alloc::vec::Vec<(usize, u16)> {
361 fn push(&mut self, value: (usize, u16)) {
362 self.push(value);
363 }
364
365 fn clear(&mut self) {
366 self.clear();
367 }
368
369 fn len(&self) -> usize {
370 self.len()
371 }
372
373 fn is_empty(&self) -> bool {
374 self.is_empty()
375 }
376
377 fn as_slice(&self) -> &[(usize, u16)] {
378 self
379 }
380 }
381
382 async fn tune_detector_manually<'a>(
383 detector: &mut DetectorDriver<MockTimeSource, MockPwm<'a>, MockAdc<'a>>,
384 buf: &mut [u16],
385 driver: &MockDriver,
386 ) {
387 let mut tuning = core::pin::pin!(detector.tune(buf));
388
389 poll_fn(|cx| match tuning.as_mut().poll(cx) {
390 core::task::Poll::Ready(_) => core::task::Poll::Ready(()),
391 core::task::Poll::Pending => {
392 cx.waker().wake_by_ref();
393 driver.advance(Duration::from_secs(2));
394
395 core::task::Poll::Pending
396 }
397 })
398 .await;
399 }
400
401 fn generate_noisy_signal<
402 I: SliceIndex<[u16], Output = [u16]> + RangeBounds<usize>,
403 S: RangeBounds<i16> + SampleRange<i16> + Clone,
404 >(
405 samples: &mut [u16],
406 amplitude: S,
407 range: I,
408 ) {
409 let mut noise = wyrand::WyRand::new(141);
410
411 let samples = samples
412 .get_mut(range)
413 .expect("Range to be sized the same or smaller than the slice");
414
415 for sample in samples.iter_mut() {
416 *sample = ((*sample as i16) - noise.random_range(amplitude.clone())) as u16;
417 }
418 }
419
420 fn generate_drop_signal(samples: &mut [u16]) {
421 // Example voltage drop signal
422 samples[5] -= 60;
423 samples[6] -= 55;
424 samples[7] -= 50;
425 samples[8] -= 45;
426 samples[9] -= 40;
427 samples[10] -= 35;
428 samples[11] -= 30;
429 samples[12] -= 28;
430 samples[13] -= 26;
431 samples[14] -= 24;
432 samples[15] -= 22;
433 samples[16] -= 20;
434 samples[17] -= 18;
435 samples[18] -= 16;
436 samples[19] -= 14;
437 samples[20] -= 12;
438 samples[21] -= 11;
439 samples[22] -= 10;
440 samples[23] -= 9;
441 samples[24] -= 8;
442 samples[25] -= 7;
443 samples[26] -= 6;
444 samples[27] -= 5;
445 samples[28] -= 4;
446 samples[29] -= 3;
447 samples[30] -= 2;
448 samples[31] -= 1;
449 }
450
451 #[pollster::test]
452 async fn tuning_cycle_completes() {
453 let driver = embassy_time::MockDriver::get();
454 driver.reset();
455 let device = MockMachine::default();
456 let pwm = MockPwm(&device);
457 let adc = MockAdc(&device);
458
459 let mut detector = DetectorDriver::new(Default::default(), MockTimeSource, pwm, adc);
460 let mut buf = [0; 16];
461
462 tune_detector_manually(&mut detector, &mut buf, driver).await;
463
464 assert_eq!(detector.state.max_duty.get(), 108);
465 assert_eq!(detector.state.duty.get(), 90);
466 assert_eq!(detector.adc.sample_average(&mut buf).await, 1260);
467 }
468
469 #[pollster::test]
470 async fn tick_updates_only_on_raised_levels() {
471 let driver = embassy_time::MockDriver::get();
472 driver.reset();
473 let device = MockMachine::default();
474 let pwm = MockPwm(&device);
475 let adc = MockAdc(&device);
476
477 let detector = DetectorDriver::new(Default::default(), MockTimeSource, pwm, adc);
478
479 detector.state.max_duty.set(98);
480 detector.state.duty.set(64);
481 detector.state.avg.set(896);
482
483 {
484 let mut tick = core::pin::pin!(detector.tick(|a| {
485 assert_eq!(
486 a,
487 DetectorUpdate::Tick {
488 timestamp: 0,
489 level: 300
490 }
491 );
492 }));
493
494 poll_fn(|cx| match tick.as_mut().poll(cx) {
495 core::task::Poll::Ready(_) => core::task::Poll::Ready(()),
496 core::task::Poll::Pending => {
497 if detector.state.warn_level.get() > 255 {
498 return core::task::Poll::Ready(());
499 }
500 detector.state.strikes.set(300);
501 driver.advance(Duration::from_secs(1));
502 cx.waker().wake_by_ref();
503 core::task::Poll::Pending
504 }
505 })
506 .await;
507 }
508
509 assert_eq!(detector.state.warn_level.get(), 553);
510 }
511
512 #[pollster::test]
513 async fn detection_updates_only_on_large_enough_signals() {
514 let driver = embassy_time::MockDriver::get();
515 driver.reset();
516 let device = MockMachine::default();
517 let pwm = MockPwm(&device);
518 let adc = MockAdc(&device);
519
520 let detector = DetectorDriver::new(Default::default(), MockTimeSource, pwm, adc);
521 detector.pwm.0.pwm.set(64);
522 detector.state.max_duty.set(98);
523 detector.state.duty.set(64);
524 detector.state.avg.set(896);
525
526 let mut samples = alloc::vec![0; BLOCK_SIZE];
527
528 detector.adc.sample(&mut samples).await;
529
530 samples[5] -= 3;
531 samples[6] -= 1;
532
533 let mut peaks = alloc::vec::Vec::with_capacity(512);
534
535 detector.detect_from_sample(0, &samples, &mut peaks, |_| {
536 panic!("This update function shouldn't be called");
537 });
538
539 assert_eq!(peaks.len(), 0);
540
541 generate_drop_signal(&mut samples);
542
543 let expected_peaks = alloc::vec![(0, 25), (8, 21)];
544 let called = Cell::new(false);
545
546 detector.detect_from_sample(0, &samples, &mut peaks, |update: DetectorUpdate<'_>| {
547 called.set(true);
548 assert_eq!(
549 update,
550 DetectorUpdate::Detection {
551 timestamp: 0,
552 average: detector.state.avg.get(),
553 samples: &samples,
554 peaks: expected_peaks.as_slice()
555 }
556 )
557 });
558
559 assert_eq!(peaks.len(), 2);
560 assert!(called.get());
561
562 called.set(false);
563
564 let mut samples = alloc::vec![0; BLOCK_SIZE];
565 let expected_peaks = alloc::vec![(16, 15), (24, 19), (32, 18)];
566
567 detector.adc.sample(&mut samples).await;
568
569 // Big noisy signal
570 generate_noisy_signal(&mut samples, -40..=40, 20..55);
571
572 detector.detect_from_sample(0, &samples, &mut peaks, |update| {
573 called.set(true);
574 assert_eq!(
575 update,
576 DetectorUpdate::Detection {
577 timestamp: 0,
578 average: detector.state.avg.get(),
579 samples: &samples,
580 peaks: expected_peaks.as_slice()
581 }
582 )
583 });
584
585 assert!(called.get());
586 assert_eq!(peaks.len(), 3);
587
588 detector.adc.sample(&mut samples).await;
589
590 // Smaller noisy signal
591 generate_noisy_signal(&mut samples, -15..=15, 10..75);
592
593 detector.detect_from_sample(0, &samples, &mut peaks, |_| {
594 panic!("This update function shouldn't be called");
595 });
596
597 assert_eq!(peaks.len(), 0);
598 }
599
600 #[pollster::test]
601 async fn detection_sensitivity_can_be_configured() {
602 let driver = embassy_time::MockDriver::get();
603 driver.reset();
604 let device = MockMachine::default();
605 let pwm = MockPwm(&device);
606 let adc = MockAdc(&device);
607
608 let detector = DetectorDriver::new(Default::default(), MockTimeSource, pwm, adc);
609 detector.pwm.0.pwm.set(64);
610 detector.state.max_duty.set(98);
611 detector.state.duty.set(64);
612 detector.state.avg.set(896);
613
614 let mut samples = alloc::vec![0; BLOCK_SIZE];
615 let mut peaks = alloc::vec::Vec::with_capacity(BLOCK_SIZE);
616
617 // Require bigger size blips.
618 detector.config.blip_size.set(3);
619
620 detector.adc.sample(&mut samples).await;
621 generate_drop_signal(&mut samples);
622
623 detector.detect_from_sample(0, &samples, &mut peaks, |_update| {
624 panic!("Update shouldn't be called");
625 });
626
627 assert_eq!(peaks.len(), 2);
628
629 // Require bigger threshold for blip detection
630 detector.config.blip_size.set(2);
631 detector.config.blip_threshold.set(24);
632
633 detector.detect_from_sample(0, &samples, &mut peaks, |_update| {
634 panic!("Update shouldn't be called");
635 });
636
637 // Update didn't call because detected blip wasn't big enough
638 assert_eq!(peaks.len(), 1);
639 }
640}