restructured to common clipping of pr<0 in "pr" and "pr2"

midas/main.py

@@ -379,66 +379,68 @@ def do_main(
         )
 
         result = separator.combine(adjusted)
-    elif adjustment_type == "pr":
-        mode = SeparationMode.MULTIPLICATIVE
-        reference = calibration_reference[0]
-        reference.data = np.clip(reference.data.core_data(), 0, None)
-        model = calibration_model[0]
-        model.data = np.clip(model.data.core_data(), 0, None)
-        adj = adjustment_model[0]
-        adj.data = np.clip(adj.data.core_data(), 0, None)
-
-        spline_kwargs = dict(smoothing_parameter_method="simple")
-        if ssr:
-            spline_kwargs["ssr"] = dict(seed=0)
-
-        statistics = DOYStatistics(
-            reference_period,
-            kernels=DOYSmootingSpline(**spline_kwargs),
-            client=scheduler.client,
-        )
-        statistics.calibrate(reference, model)
-        result = statistics.adjust(adj)
-        result_data = result.core_data()
-        result.data = np.where(result_data < pr_low_threshold, 0.0, result_data)
-    elif adjustment_type == "pr2":
-        mode = SeparationMode.MULTIPLICATIVE
-        separator = TimeSeparator(mode=mode)
-        calibration_reference[0].data = np.clip(calibration_reference[0].core_data(), 0, None)
-        reference = separator.separate(calibration_reference[0])
-        calibration_model[0].data = np.clip(calibration_model[0].core_data(), 0, None)
-        model = separator.separate(calibration_model[0])
-        adjustment_model[0].data = np.clip(adjustment_model[0].core_data(), 0, None)
-        adj = separator.separate(adjustment_model[0])
-
-        anomaly_statistics = DOYStatistics(
-            reference_period,
-            kernels=DOYSmootingSpline("simple"),
-            client=scheduler.client,
-        )
-        anomaly_statistics.calibrate(
-            reference[TimeSeparator.DAILY_ANOMALIES],
-            model[TimeSeparator.DAILY_ANOMALIES],
-        )
+    elif adjustment_type in ["pr", "pr2"]:
+        calibration_reference[0].data = np.where(
+            calibration_reference[0].core_data() < 0, 0,
+            calibration_reference[0].core_data())
+        calibration_model[0].data = np.where(
+            calibration_model[0].core_data() < 0, 0,
+            calibration_model[0].core_data())
+        adjustment_model[0].data = np.where(
+            adjustment_model[0].core_data() < 0, 0,
+            adjustment_model[0].core_data())
+        if adjustment_type == "pr":
+            mode = SeparationMode.MULTIPLICATIVE
+            reference = calibration_reference[0]
+            model = calibration_model[0]
+            adj = adjustment_model[0]
+
+            spline_kwargs = dict(smoothing_parameter_method="simple")
+            if ssr:
+                spline_kwargs["ssr"] = dict(seed=0)
+
+            statistics = DOYStatistics(
+                reference_period,
+                kernels=DOYSmootingSpline(**spline_kwargs),
+                client=scheduler.client,
+            )
+            statistics.calibrate(reference, model)
+            result = statistics.adjust(adj)
+        elif adjustment_type == "pr2":
+            mode = SeparationMode.MULTIPLICATIVE
+            separator = TimeSeparator(mode=mode)
+            reference = separator.separate(calibration_reference[0])
+            model = separator.separate(calibration_model[0])
+            adj = separator.separate(adjustment_model[0])
+
+            anomaly_statistics = DOYStatistics(
+                reference_period,
+                kernels=DOYSmootingSpline("simple"),
+                client=scheduler.client,
+            )
+            anomaly_statistics.calibrate(
+                reference[TimeSeparator.DAILY_ANOMALIES],
+                model[TimeSeparator.DAILY_ANOMALIES],
+            )
 
-        adjusted = {}
-        adjusted[TimeSeparator.DAILY_ANOMALIES] = anomaly_statistics.adjust(
-            adj[TimeSeparator.DAILY_ANOMALIES]
-        )
+            adjusted = {}
+            adjusted[TimeSeparator.DAILY_ANOMALIES] = anomaly_statistics.adjust(
+                adj[TimeSeparator.DAILY_ANOMALIES]
+            )
 
-        multiday_statistics = DOYStatistics(
-            reference_period,
-            kernels=DOYSmootingSpline("simple"),
-            client=scheduler.client,
-        )
-        multiday_statistics.calibrate(
-            reference[TimeSeparator.MULTI_DAY], model[TimeSeparator.MULTI_DAY]
-        )
-        adjusted[TimeSeparator.MULTI_DAY] = multiday_statistics.adjust(
-            adj[TimeSeparator.MULTI_DAY]
-        )
+            multiday_statistics = DOYStatistics(
+                reference_period,
+                kernels=DOYSmootingSpline("simple"),
+                client=scheduler.client,
+            )
+            multiday_statistics.calibrate(
+                reference[TimeSeparator.MULTI_DAY], model[TimeSeparator.MULTI_DAY]
+            )
+            adjusted[TimeSeparator.MULTI_DAY] = multiday_statistics.adjust(
+                adj[TimeSeparator.MULTI_DAY]
+            )
+            result = separator.combine(adjusted)
 
-        result = separator.combine(adjusted)
         result_data = result.core_data()
         result.data = np.where(result_data < pr_low_threshold, 0.0, result_data)
     else: