SWI calculation now may take wind speed criterion into accountHEAD0.1.0master

1 files changed, 16 insertions, 3 deletions

diff --git a/szilagyi/__init__.py b/szilagyi/__init__.py
index bfde154..70c073e 100644
--- a/szilagyi/__init__.py
+++ b/szilagyi/__init__.py
@@ -13,15 +13,23 @@ There is only one function provided, so the use should be straight-forward. For
 	-4.88
 """
 
-from pint.quantity import Quantity
+from pint import quantity as pint_
 
 from ._nomogram import calculate_swi as _calculate_swi
 from ._nomogram import MAX_TEMPERATURE_DIFFERENCE, MAX_CONVECTIVE_CLOUD_DEPTH
 
 
-def calculate_swi(temperature_difference: Quantity, convective_cloud_depth: Quantity) -> float:
+WIND_SPEED_LIMIT = 40
+
+
+def calculate_swi(
+	temperature_difference: pint_.Quantity,
+	convective_cloud_depth: pint_.Quantity,
+	wind_speed_at_850hpa: pint_.Quantity = None
+) -> float:
 	"""
-	Calculates the Szilagyi Waterspout Index for *temperature_difference* and *convective_cloud_depth*.
+	Calculates the Szilagyi Waterspout Index for *temperature_difference* and *convective_cloud_depth*. Optionally,
+	*wind_speed_at_850hpa* can be provided to filter out cases by this parameter.
 
 	Function will raise ValueError if input values are not convertible to their expected units or if they are not in range
 	of the original nomogram. The maximum values are available as MAX_TEMPERATURE_DIFFERENCE and MAX_CONVECTIVE_CLOUD_DEPTH
@@ -31,6 +39,11 @@ def calculate_swi(temperature_difference: Quantity, convective_cloud_depth: Quan
 		raise ValueError("temperature_difference must be convertible to Celsius degrees")
 	if not convective_cloud_depth.is_compatible_with("ft"):
 		raise ValueError("convective_cloud_depth must be convertible to Feet")
+	if wind_speed_at_850hpa:
+		if not wind_speed_at_850hpa.is_compatible_with("kts"):
+			raise ValueError("wind_speed_at_850hpa must be convertible to Knots")
+		if wind_speed_at_850hpa.m_as("kts") > WIND_SPEED_LIMIT:
+			return -10
 	temperature_difference = temperature_difference.m_as("delta_degC")
 	convective_cloud_depth = convective_cloud_depth.m_as("ft")
 	if temperature_difference < 0 or temperature_difference > MAX_TEMPERATURE_DIFFERENCE: