Change if to where

policyengine_uk/variables/gov/hmrc/tax_free_childcare/tax_free_childcare_benefits.py

@@ -41,23 +41,21 @@ def formula(benunit, period, parameters):
             meets_age_condition & meets_income_condition & childcare_eligible
         )
 
-        # Determine the maximum eligible childcare cost for a single child
-        max_amount = 0
+        # Calculate maximum eligible childcare cost using vectorized operations
         child_mask = benunit.members("is_child", period)
-        for member in benunit.members:
-            disabled = benunit.members("is_disabled", period)[member]
-            max_amount = where(
-                child_mask[member],
+        disabled = benunit.members("is_disabled", period)
+        max_amount = where(
+            child_mask,
+            where(
+                is_eligible,
                 where(
-                    is_eligible,
-                    where(
-                        disabled,
-                        p.disabled_child.values,
-                        p.standard_child.values,
-                    ),
-                    max_amount,
+                    disabled,
+                    p.disabled_child.values,
+                    p.standard_child.values,
                 ),
-                max_amount,
-            )
+                0,
+            ),
+            0,
+        ).max(axis=0)
 
         return where(is_eligible, max_amount, 0)