Source code for pelagos_py.steps.input_output.correct_values

# This file is part of pelagos_py.
#
# Copyright 2025-2026 National Oceanography Centre and The Contributors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Pipeline step for applying a linear correction (slope/intercept) to a variable.

Generic enough for sensor cross-calibration (slope + intercept), unit conversions,
or any other affine rescaling. An optional
``expected_range`` makes the correction self-detecting: it is applied only when
the data looks like it still needs it, so the same config keeps working even after
upstream input files are fixed.
"""

#### Mandatory imports ####
from pelagos_py.steps.base_step import BaseStep, register_step

#### Custom imports ####
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
from pelagos_py.utils import fig_spec


@register_step
[docs] class CorrectValues(BaseStep): """ Apply an affine correction ``corrected = slope * value + intercept`` to a variable. The correction is conditional when ``expected_range`` is given: the median of the valid (non-NaN) data is compared against ``[min, max]``, and the correction is applied only when the median falls *outside* that range (i.e. the data still looks uncorrected). When ``expected_range`` is omitted the correction is always applied. This makes a config robust to upstream fixes: a unit conversion targeting an ``expected_range`` quietly skips files that already arrive in the right units. (CNDC does not need this: Derive CTD reads the units attribute and converts S/m -> mS/cm for gsw itself.) Parameters ---------- target_variable : str Name of the variable to correct (e.g. ``"CNDC"``). output_as : str or list, optional Name(s) to write the result under. Defaults to ``target_variable`` (in-place). Set it to a new name to copy/rename, e.g. ``LATITUDE_GPS`` -> ``LATITUDE`` with ``slope: 1`` leaves the values unchanged and just exposes the new name. A list writes the same result under several names, e.g. ``[CHLA, CHLA_MID]`` keeps a mid-processing snapshot in ``CHLA_MID`` while later steps continue working on ``CHLA``. slope : float, optional Multiplicative factor (default ``1.0``). For a x10 unit conversion, set ``10``. intercept : float, optional Additive offset applied after scaling (default ``0.0``). Use for alignment. expected_range : list, optional ``[min, max]`` for the *corrected* variable. The correction is applied only when the data's median falls outside this range. If omitted, the correction is always applied. time_start, time_end : str, optional Restrict the correction to a TIME window (e.g. ``"2024-07-01T00:00:00"``). Points outside the window keep their raw value. Either bound may be omitted. corrected_units : str, optional Units string written to the output variable's attributes after a correction is applied (e.g. ``"mS/cm"``). Left untouched if omitted or if no correction runs. append_description, overwrite_description : str, optional Note written to the output variable's ``comment`` attribute (e.g. ``"Renamed from LATITUDE_GPS"``). ``append_description`` adds to any existing comment; ``overwrite_description`` replaces it. Set at most one. Examples -------- Example usage in a pipeline configuration: .. code-block:: yaml steps: - name: Correct Values parameters: target_variable: TEMP slope: 1.001 intercept: -0.05 expected_range: [-2.5, 40] diagnostics: false """ step_name = "Correct Values" required_variables = [] provided_variables = [] parameter_schema = { "target_variable": { "type": str, "required": True, "description": "Name of the variable to correct (e.g. 'CNDC').", }, "output_as": { "type": [str, list], "default": None, "description": "Name(s) to write the result under (default: target_variable, " "i.e. in place). A new name copies/renames; a list writes the " "same result under several names (e.g. [CHLA, CHLA_MID]).", }, "slope": { "type": float, "default": 1.0, "description": "Multiplicative factor (corrected = slope * value + intercept). " "For a simple x10 unit conversion, set 10.", }, "intercept": { "type": float, "default": 0.0, "description": "Additive offset applied after scaling (corrected = slope * value + intercept). " "Use for sensor alignment.", }, "expected_range": { "type": list, "default": None, "description": "Optional [min, max] for the corrected variable. The correction is applied " "only when the data's median falls OUTSIDE this range. If omitted, the " "correction is always applied.", }, "corrected_units": { "type": str, "default": None, "description": "Optional units string written to the output variable's attributes " "after a correction is applied (e.g. 'mS/cm').", }, "time_start": { "type": str, "default": None, "description": "Optional ISO timestamp; only points at/after this TIME are corrected.", }, "time_end": { "type": str, "default": None, "description": "Optional ISO timestamp; only points at/before this TIME are corrected.", }, "append_description": { "type": str, "default": None, "description": "Optional note appended to the output variable's 'comment' attribute.", }, "overwrite_description": { "type": str, "default": None, "description": "Optional note that replaces the output variable's 'comment' attribute.", }, } def run(self): self.check_data() self.data = self.context["data"] var = self.target_variable if var not in self.data: raise ValueError( f"[{self.name}] target_variable '{var}' not found in dataset. " f"Available variables: {list(self.data.data_vars)}." ) outs = ( list(self.output_as) if isinstance(self.output_as, (list, tuple)) else [self.output_as or var] ) if self.append_description is not None and self.overwrite_description is not None: raise ValueError( f"[{self.name}] set only one of 'append_description' / 'overwrite_description'." ) vals = self.data[var].values.astype(float) self._raw_data = vals.copy() self.applied = False valid_mask = ~np.isnan(vals) if not np.any(valid_mask): self.log_warn(f"'{var}' has no valid (non-NaN) values; nothing to correct.") self.context["data"] = self.data return self.context # Restrict the correction to a TIME window when requested. window = np.ones(vals.shape, dtype=bool) if self.time_start is not None or self.time_end is not None: if "TIME" not in self.data: raise ValueError( f"[{self.name}] time_start/time_end need a 'TIME' variable in the data." ) times = self.data["TIME"].values if self.time_start is not None: window &= times >= np.datetime64(self.time_start) if self.time_end is not None: window &= times <= np.datetime64(self.time_end) # Decide whether the correction is needed, judging by the windowed data. do_scale = True if self.expected_range is not None: lo, hi = float(self.expected_range[0]), float(self.expected_range[1]) sample = vals[valid_mask & window] median_val = float(np.nanmedian(sample)) if sample.size else np.nan if np.isfinite(median_val) and lo <= median_val <= hi: self.log( f"'{var}' median ({median_val:.4g}) is within expected range " f"[{lo}, {hi}]; skipping correction." ) do_scale = False elif np.isfinite(median_val): self.log( f"'{var}' median ({median_val:.4g}) is outside expected range " f"[{lo}, {hi}]; applying correction." ) # Nothing changes when there is no scaling, no rename/copy and no comment to set. no_description = self.append_description is None and self.overwrite_description is None if not do_scale and outs == [var] and no_description: self.context["data"] = self.data return self.context # Apply the affine correction within the window (NaNs propagate harmlessly). corrected = vals.copy() if do_scale: corrected[window] = self.slope * vals[window] + self.intercept self.applied = True # Write to each output name (a copy/rename when it differs from target_variable). for out in outs: self.data[out] = self.data[var].copy(data=corrected) self._outs = outs names = ", ".join(f"'{o}'" for o in outs) if self.applied: self.log( f"Applied correction to {names}: corrected = {self.slope} * value + {self.intercept}." ) copies = [o for o in outs if o != var] if copies: self.log(f"Wrote '{var}' to {', '.join(repr(o) for o in copies)}.") for out in outs: if self.corrected_units is not None: self.data[out].attrs["units"] = self.corrected_units if self.overwrite_description is not None: self.data[out].attrs["comment"] = self.overwrite_description elif self.append_description is not None: existing = self.data[out].attrs.get("comment", "") self.data[out].attrs["comment"] = ( f"{existing} {self.append_description}".strip() if existing else self.append_description ) if self.corrected_units is not None: self.log(f"Set {names} units to '{self.corrected_units}'.") if self.diagnostics: self.plot_diagnostics() self.context["data"] = self.data return self.context def plot_diagnostics(self): if not self.applied: return var = self.target_variable corrected = self.data[self._outs[0]].values # Plot against TIME if available, otherwise against sample index. if "TIME" in self.data: x = self.data["TIME"].values xlabel = "Time" else: x = np.arange(len(corrected)) xlabel = "Sample index" matplotlib.use("tkagg") fig, axes = fig_spec.new_fig() ax = axes[0][0] fig_spec.points(ax, x, self._raw_data, color=fig_spec.FLAGGED, label="Raw") fig_spec.points(ax, x, corrected, color=fig_spec.CATEGORY[1], label="Corrected") if self.expected_range is not None: lo, hi = float(self.expected_range[0]), float(self.expected_range[1]) ax.axhline(hi, color="black", linestyle="--", alpha=0.6, linewidth=1, label=f"Max ({hi})") ax.axhline(lo, color="black", linestyle="--", alpha=0.6, linewidth=1, label=f"Min ({lo})") if xlabel == "Time": fig_spec.date_axis(ax, which="x") ylabel = fig_spec.axis_label(var, self.data[var].attrs.get("units")) fig_spec.style_axes(ax, xlabel=xlabel, ylabel=ylabel) fig_spec.legend(ax) fig_spec.finish( fig, suptitle=f"Value Correction: {var}\n" f"(corrected = {self.slope} * value + {self.intercept})", ) plt.show(block=True)