[aqi] Add extended_range option for over-range AQI values (#17570)

Co-authored-by: jas <jas@asspa.in>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

esphome/components/aqi/__init__.py

@@ -7,6 +7,7 @@ AQICalculatorType = aqi_ns.enum("AQICalculatorType")
 
 CONF_AQI = "aqi"
 CONF_CALCULATION_TYPE = "calculation_type"
+CONF_EXTENDED_RANGE = "extended_range"
 
 AQI_CALCULATION_TYPE = {
     "CAQI": AQICalculatorType.CAQI_TYPE,

esphome/components/aqi/abstract_aqi_calculator.h

@@ -6,7 +6,7 @@ namespace esphome::aqi {
 
 class AbstractAQICalculator {
  public:
-  virtual uint16_t get_aqi(float pm2_5_value, float pm10_0_value) = 0;
+  virtual uint16_t get_aqi(float pm2_5_value, float pm10_0_value, bool extended_range) = 0;
 };
 
 }  // namespace esphome::aqi

esphome/components/aqi/aqi_calculator.h

@@ -11,10 +11,12 @@ namespace esphome::aqi {
 
 class AQICalculator : public AbstractAQICalculator {
  public:
-  uint16_t get_aqi(float pm2_5_value, float pm10_0_value) override {
-    float pm2_5_index = calculate_index(pm2_5_value, PM2_5_GRID);
-    float pm10_0_index = calculate_index(pm10_0_value, PM10_0_GRID);
+  uint16_t get_aqi(float pm2_5_value, float pm10_0_value, bool extended_range) override {
+    float pm2_5_index = calculate_index(pm2_5_value, PM2_5_GRID, extended_range);
+    float pm10_0_index = calculate_index(pm10_0_value, PM10_0_GRID, extended_range);
     float aqi = std::max({pm2_5_index, pm10_0_index, 0.0f});
+    // extended_range lets the index run past the standard maximum, so clamp to the sensor's range.
+    aqi = std::min(aqi, static_cast<float>(std::numeric_limits<uint16_t>::max()));
     return static_cast<uint16_t>(std::lround(aqi));
   }
 
@@ -30,7 +32,7 @@ class AQICalculator : public AbstractAQICalculator {
       {35.5f, 55.5f},
       {55.5f, 125.5f},
       {125.5f, 225.5f},
-      {225.5f, std::numeric_limits<float>::max()}
+      {225.5f, 500.4f}  // EPA 2024: AQI 301-500 maps to PM2.5 225.5-500.4 ug/m3
       // clang-format on
   };
 
@@ -41,11 +43,11 @@ class AQICalculator : public AbstractAQICalculator {
       {155.0f, 255.0f},
       {255.0f, 355.0f},
       {355.0f, 425.0f},
-      {425.0f, std::numeric_limits<float>::max()}
+      {425.0f, 604.0f}  // EPA: AQI 301-500 maps to PM10 425-604 ug/m3 (top of the 401-500 band)
       // clang-format on
   };
 
-  static float calculate_index(float value, const float array[NUM_LEVELS][2]) {
+  static float calculate_index(float value, const float array[NUM_LEVELS][2], bool extended_range) {
     int grid_index = get_grid_index(value, array);
     if (grid_index == -1) {
       return -1.0f;
@@ -55,14 +57,22 @@ class AQICalculator : public AbstractAQICalculator {
     float conc_lo = array[grid_index][0];
     float conc_hi = array[grid_index][1];
 
-    return (value - conc_lo) * (aqi_hi - aqi_lo) / (conc_hi - conc_lo) + aqi_lo;
+    float index = (value - conc_lo) * (aqi_hi - aqi_lo) / (conc_hi - conc_lo) + aqi_lo;
+
+    // Concentrations above the highest breakpoint run the linear fit past aqi_hi. By default we
+    // clamp to the standard maximum; with extended_range we keep the extrapolated "over-range"
+    // value so heavy pollution reports numbers beyond what the standard defines.
+    if (grid_index == NUM_LEVELS - 1 && !extended_range && index > aqi_hi) {
+      return aqi_hi;
+    }
+    return index;
   }
 
   static int get_grid_index(float value, const float array[NUM_LEVELS][2]) {
     for (int i = 0; i < NUM_LEVELS; i++) {
-      const bool in_range =
-          (value >= array[i][0]) && ((i == NUM_LEVELS - 1) ? (value <= array[i][1])   // last bucket inclusive
-                                                           : (value < array[i][1]));  // others exclusive on hi
+      // The top band is open-ended: any value at or above its lower breakpoint falls into it,
+      // and calculate_index() decides whether to clamp or extrapolate.
+      const bool in_range = (value >= array[i][0]) && (i == NUM_LEVELS - 1 || value < array[i][1]);
       if (in_range) {
         return i;
       }

esphome/components/aqi/aqi_sensor.cpp

@@ -24,6 +24,7 @@ void AQISensor::setup() {
 void AQISensor::dump_config() {
   ESP_LOGCONFIG(TAG, "AQI Sensor:");
   ESP_LOGCONFIG(TAG, "  Calculation Type: %s", this->aqi_calc_type_ == AQI_TYPE ? "AQI" : "CAQI");
+  ESP_LOGCONFIG(TAG, "  Extended Range: %s", this->extended_range_ ? "enabled" : "disabled");
   if (this->pm_2_5_sensor_ != nullptr) {
     ESP_LOGCONFIG(TAG, "  PM2.5 Sensor: '%s'", this->pm_2_5_sensor_->get_name().c_str());
   }
@@ -44,7 +45,7 @@ void AQISensor::calculate_aqi_() {
     return;
   }
 
-  uint16_t aqi = calculator->get_aqi(this->pm_2_5_value_, this->pm_10_0_value_);
+  uint16_t aqi = calculator->get_aqi(this->pm_2_5_value_, this->pm_10_0_value_, this->extended_range_);
   this->publish_state(aqi);
 }
 

esphome/components/aqi/aqi_sensor.h

@@ -14,6 +14,7 @@ class AQISensor final : public sensor::Sensor, public Component {
   void set_pm_2_5_sensor(sensor::Sensor *sensor) { this->pm_2_5_sensor_ = sensor; }
   void set_pm_10_0_sensor(sensor::Sensor *sensor) { this->pm_10_0_sensor_ = sensor; }
   void set_aqi_calculation_type(AQICalculatorType type) { this->aqi_calc_type_ = type; }
+  void set_extended_range(bool extended_range) { this->extended_range_ = extended_range; }
 
  protected:
   void calculate_aqi_();
@@ -21,6 +22,7 @@ class AQISensor final : public sensor::Sensor, public Component {
   sensor::Sensor *pm_2_5_sensor_{nullptr};
   sensor::Sensor *pm_10_0_sensor_{nullptr};
   AQICalculatorType aqi_calc_type_{AQI_TYPE};
+  bool extended_range_{false};
   AQICalculatorFactory aqi_calculator_factory_;
 
   float pm_2_5_value_{NAN};

esphome/components/aqi/caqi_calculator.h

@@ -9,25 +9,28 @@ namespace esphome::aqi {
 
 class CAQICalculator : public AbstractAQICalculator {
  public:
-  uint16_t get_aqi(float pm2_5_value, float pm10_0_value) override {
+  // The CAQI (CITEAIR) scale defines no maximum: its top "Very high" class is simply ">100". We
+  // therefore always extrapolate the top band past 100 without limit, so the extended_range flag
+  // (which lifts the AQI calculator's fixed 500 cap) has no meaning here and is ignored.
+  uint16_t get_aqi(float pm2_5_value, float pm10_0_value, bool /*extended_range*/) override {
     float pm2_5_index = calculate_index(pm2_5_value, PM2_5_GRID);
     float pm10_0_index = calculate_index(pm10_0_value, PM10_0_GRID);
     float aqi = std::max({pm2_5_index, pm10_0_index, 0.0f});
+    aqi = std::min(aqi, static_cast<float>(std::numeric_limits<uint16_t>::max()));
     return static_cast<uint16_t>(std::lround(aqi));
   }
 
  protected:
-  static constexpr int NUM_LEVELS = 5;
+  static constexpr int NUM_LEVELS = 4;
 
-  static constexpr int INDEX_GRID[NUM_LEVELS][2] = {{0, 25}, {26, 50}, {51, 75}, {76, 100}, {101, 400}};
+  static constexpr int INDEX_GRID[NUM_LEVELS][2] = {{0, 25}, {26, 50}, {51, 75}, {76, 100}};
 
   static constexpr float PM2_5_GRID[NUM_LEVELS][2] = {
       // clang-format off
       {0.0f, 15.1f},
       {15.1f, 30.1f},
       {30.1f, 55.1f},
-      {55.1f, 110.1f},
-      {110.1f, std::numeric_limits<float>::max()}
+      {55.1f, 110.1f}
       // clang-format on
   };
 
@@ -36,8 +39,7 @@ class CAQICalculator : public AbstractAQICalculator {
       {0.0f, 25.1f},
       {25.1f, 50.1f},
       {50.1f, 90.1f},
-      {90.1f, 180.1f},
-      {180.1f, std::numeric_limits<float>::max()}
+      {90.1f, 180.1f}
       // clang-format on
   };
 
@@ -52,14 +54,15 @@ class CAQICalculator : public AbstractAQICalculator {
     float conc_lo = array[grid_index][0];
     float conc_hi = array[grid_index][1];
 
+    // The top band is open-ended (see get_grid_index), so for concentrations above the last
+    // breakpoint this linear fit extrapolates past 100 unbounded, matching CAQI's open ">100" class.
     return (value - conc_lo) * (aqi_hi - aqi_lo) / (conc_hi - conc_lo) + aqi_lo;
   }
 
   static int get_grid_index(float value, const float array[NUM_LEVELS][2]) {
     for (int i = 0; i < NUM_LEVELS; i++) {
-      const bool in_range =
-          (value >= array[i][0]) && ((i == NUM_LEVELS - 1) ? (value <= array[i][1])   // last bucket inclusive
-                                                           : (value < array[i][1]));  // others exclusive on hi
+      // The top band is open-ended: any value at or above its lower breakpoint falls into it.
+      const bool in_range = (value >= array[i][0]) && (i == NUM_LEVELS - 1 || value < array[i][1]);
       if (in_range) {
         return i;
       }

esphome/components/aqi/sensor.py

@@ -8,14 +8,25 @@ from esphome.const import (
     STATE_CLASS_MEASUREMENT,
 )
 
-from . import AQI_CALCULATION_TYPE, CONF_CALCULATION_TYPE, aqi_ns
+from . import AQI_CALCULATION_TYPE, CONF_CALCULATION_TYPE, CONF_EXTENDED_RANGE, aqi_ns
 
 CODEOWNERS = ["@jasstrong"]
 DEPENDENCIES = ["sensor"]
 
 AQISensor = aqi_ns.class_("AQISensor", sensor.Sensor, cg.Component)
 
-CONFIG_SCHEMA = (
+
+def _validate_extended_range(config):
+    if CONF_EXTENDED_RANGE in config and config[CONF_CALCULATION_TYPE] == "CAQI":
+        raise cv.Invalid(
+            f"'{CONF_EXTENDED_RANGE}' is not supported with 'calculation_type: CAQI'. "
+            "CAQI has no maximum value by specification, so it is always reported unbounded.",
+            [CONF_EXTENDED_RANGE],
+        )
+    return config
+
+
+CONFIG_SCHEMA = cv.All(
     sensor.sensor_schema(
         AQISensor,
         accuracy_decimals=0,
@@ -29,9 +40,11 @@ CONFIG_SCHEMA = (
             cv.Required(CONF_CALCULATION_TYPE): cv.enum(
                 AQI_CALCULATION_TYPE, upper=True
             ),
+            cv.Optional(CONF_EXTENDED_RANGE): cv.boolean,
         }
     )
-    .extend(cv.COMPONENT_SCHEMA)
+    .extend(cv.COMPONENT_SCHEMA),
+    _validate_extended_range,
 )
 
 
@@ -46,3 +59,4 @@ async def to_code(config):
     cg.add(var.set_pm_10_0_sensor(pm_10_0_sensor))
 
     cg.add(var.set_aqi_calculation_type(config[CONF_CALCULATION_TYPE]))
+    cg.add(var.set_extended_range(config.get(CONF_EXTENDED_RANGE, False)))

esphome/components/hm3301/hm3301.cpp

@@ -61,7 +61,7 @@ void HM3301Component::update() {
   int16_t aqi_value = -1;
   if (this->aqi_sensor_ != nullptr && pm_2_5_value != -1 && pm_10_0_value != -1) {
     aqi::AbstractAQICalculator *calculator = this->aqi_calculator_factory_.get_calculator(this->aqi_calc_type_);
-    aqi_value = calculator->get_aqi(pm_2_5_value, pm_10_0_value);
+    aqi_value = calculator->get_aqi(pm_2_5_value, pm_10_0_value, /*extended_range=*/false);
   }
 
   if (pm_1_0_value != -1) {

tests/component_tests/aqi/test_aqi.py

@@ -0,0 +1,35 @@
+"""Config-validation tests for the aqi sensor component."""
+
+import pytest
+from voluptuous import Invalid
+
+from esphome.components.aqi import CONF_CALCULATION_TYPE, CONF_EXTENDED_RANGE
+from esphome.components.aqi.sensor import _validate_extended_range
+
+
+def test_extended_range_rejected_with_caqi():
+    """extended_range has no meaning for CAQI (no spec maximum) and must be rejected."""
+    with pytest.raises(Invalid, match="CAQI"):
+        _validate_extended_range(
+            {CONF_CALCULATION_TYPE: "CAQI", CONF_EXTENDED_RANGE: True}
+        )
+
+
+def test_extended_range_rejected_with_caqi_even_when_false():
+    """The option is not allowed at all with CAQI, regardless of its value."""
+    with pytest.raises(Invalid, match="CAQI"):
+        _validate_extended_range(
+            {CONF_CALCULATION_TYPE: "CAQI", CONF_EXTENDED_RANGE: False}
+        )
+
+
+def test_extended_range_allowed_with_aqi():
+    """extended_range is valid for the US AQI calculation."""
+    config = {CONF_CALCULATION_TYPE: "AQI", CONF_EXTENDED_RANGE: True}
+    assert _validate_extended_range(config) is config
+
+
+def test_caqi_without_extended_range_ok():
+    """CAQI is fine as long as extended_range is not set."""
+    config = {CONF_CALCULATION_TYPE: "CAQI"}
+    assert _validate_extended_range(config) is config

tests/components/aqi/benchmark.yaml

@@ -0,0 +1,16 @@
+# Declares the component graph the C++ unit test build needs so that the aqi
+# component's sources (which include sensor.h) compile. to_code is suppressed by
+# the test harness; this only pulls the sensor + aqi source/include paths in.
+# Loaded with plain yaml.safe_load, so avoid lambdas / ESPHome-tagged values here.
+sensor:
+  - platform: template
+    id: pm25_sensor
+    name: "PM2.5"
+  - platform: template
+    id: pm10_sensor
+    name: "PM10"
+  - platform: aqi
+    name: "AQI"
+    pm_2_5: pm25_sensor
+    pm_10_0: pm10_sensor
+    calculation_type: AQI

tests/components/aqi/common.yaml

@@ -20,3 +20,10 @@ sensor:
     pm_2_5: pm25_sensor
     pm_10_0: pm10_sensor
     calculation_type: CAQI
+
+  - platform: aqi
+    name: "Air Quality Index (AQI, extended)"
+    pm_2_5: pm25_sensor
+    pm_10_0: pm10_sensor
+    calculation_type: AQI
+    extended_range: true

tests/components/aqi/test_aqi_calculator.cpp

@@ -0,0 +1,85 @@
+#include <gtest/gtest.h>
+
+#include "esphome/components/aqi/aqi_calculator.h"
+#include "esphome/components/aqi/caqi_calculator.h"
+
+namespace esphome::aqi::testing {
+
+// US AQI (EPA 2024): PM2.5 225.5-500.4 -> 301-500, PM10 425-604 -> 301-500.
+
+TEST(USAQI, LowRangeUnaffectedByExtendedFlag) {
+  AQICalculator calc;
+  // PM2.5 25 drives over PM10 50; well below the top band, so the flag changes nothing.
+  EXPECT_EQ(calc.get_aqi(25.0f, 50.0f, false), 81);
+  EXPECT_EQ(calc.get_aqi(25.0f, 50.0f, true), 81);
+}
+
+TEST(USAQI, HazardousInterpolatesNotPinnedAt301) {
+  AQICalculator calc;
+  // Regression guard: the old FLT_MAX top bucket collapsed every hazardous reading to 301.
+  EXPECT_EQ(calc.get_aqi(225.5f, 0.0f, false), 301);  // band start
+  EXPECT_EQ(calc.get_aqi(250.0f, 0.0f, false), 319);  // interpolated, not 301
+  EXPECT_EQ(calc.get_aqi(500.4f, 0.0f, false), 500);  // band top
+}
+
+TEST(USAQI, DefaultClampsAtStandardMaximum) {
+  AQICalculator calc;
+  EXPECT_EQ(calc.get_aqi(600.0f, 0.0f, false), 500);
+  EXPECT_EQ(calc.get_aqi(1000.0f, 0.0f, false), 500);
+  EXPECT_EQ(calc.get_aqi(0.0f, 604.0f, false), 500);  // PM10 top breakpoint
+}
+
+TEST(USAQI, ExtendedRangeExtrapolatesBeyond500) {
+  AQICalculator calc;
+  EXPECT_EQ(calc.get_aqi(600.0f, 0.0f, true), 572);
+  EXPECT_EQ(calc.get_aqi(1000.0f, 0.0f, true), 862);
+  EXPECT_EQ(calc.get_aqi(0.0f, 700.0f, true), 607);  // PM10 extrapolated past 500
+}
+
+TEST(USAQI, ExtendedRangeSaturatesUint16NoWraparound) {
+  AQICalculator calc;
+  // An absurd concentration would overflow uint16_t; it must saturate, not wrap to a small value.
+  EXPECT_EQ(calc.get_aqi(100000.0f, 0.0f, true), 65535);
+}
+
+TEST(USAQI, WorseOfTwoPollutantsWins) {
+  AQICalculator calc;
+  // PM10 604 -> 500 dominates PM2.5 25 -> 81.
+  EXPECT_EQ(calc.get_aqi(25.0f, 604.0f, false), 500);
+}
+
+// CAQI (CITEAIR): no maximum by spec -- the top ">100" class is open, so it is always unbounded
+// and the extended_range flag does not apply.
+
+TEST(CAQI, LowRange) {
+  CAQICalculator calc;
+  EXPECT_EQ(calc.get_aqi(25.0f, 50.0f, false), 50);
+}
+
+TEST(CAQI, ContinuousAt100NoPinAt101) {
+  CAQICalculator calc;
+  // Old code pinned everything above the top breakpoint to 101; now it reaches exactly 100.
+  EXPECT_EQ(calc.get_aqi(110.1f, 0.0f, false), 100);
+}
+
+TEST(CAQI, UnboundedAboveTopBand) {
+  CAQICalculator calc;
+  EXPECT_EQ(calc.get_aqi(200.0f, 0.0f, false), 139);
+  EXPECT_EQ(calc.get_aqi(2000.0f, 0.0f, false), 925);
+}
+
+TEST(CAQI, ExtendedRangeFlagIsIgnored) {
+  CAQICalculator calc;
+  // CAQI is always unbounded, so the flag must make no difference either way.
+  EXPECT_EQ(calc.get_aqi(200.0f, 0.0f, true), calc.get_aqi(200.0f, 0.0f, false));
+  EXPECT_EQ(calc.get_aqi(2000.0f, 0.0f, true), calc.get_aqi(2000.0f, 0.0f, false));
+}
+
+TEST(CAQI, SaturatesUint16NoWraparound) {
+  CAQICalculator calc;
+  // CAQI is unbounded, so an extreme reading can extrapolate past uint16_t; it must saturate,
+  // not wrap around to a small (falsely "good") value.
+  EXPECT_EQ(calc.get_aqi(200000.0f, 0.0f, false), 65535);
+}
+
+}  // namespace esphome::aqi::testing