pyelq 1.1.0__py3-none-any.whl

pyelq/component/error_model.py

@@ -0,0 +1,327 @@
+# SPDX-FileCopyrightText: 2024 Shell Global Solutions International B.V. All Rights Reserved.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+# -*- coding: utf-8 -*-
+"""Error model module."""
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Union
+
+import numpy as np
+from openmcmc import parameter
+from openmcmc.distribution.distribution import Gamma
+from openmcmc.model import Model
+from openmcmc.sampler.sampler import NormalGamma
+
+from pyelq.component.component import Component
+from pyelq.gas_species import GasSpecies
+from pyelq.meteorology import MeteorologyGroup
+from pyelq.sensor.sensor import Sensor, SensorGroup
+
+if TYPE_CHECKING:
+    from pyelq.plotting.plot import Plot
+
+
+@dataclass
+class ErrorModel(Component):
+    """Measurement precision model component for the model.
+
+    Attributes:
+        n_sensor (int): number of sensors in the sensor object used for analysis.
+        precision_index (np.ndarray): index mapping precision parameters onto observations. Will be set up differently
+            for different model types.
+        precision_parameter (parameter.Parameter): parameter object which constructs the full measurement error
+            precision matrix from the components stored in state. Will be passed to the distribution for the observed
+            when the full model is constructed.
+        prior_precision_shape (Union[np.ndarray, float]): prior shape parameters for the precision model. Set up
+            differently per model type.
+        prior_precision_rate (Union[np.ndarray, float]): prior rate parameters for the precision model. Set up
+            differently per model type.
+        initial_precision (Union[np.ndarray, float]): initial value for the precision to be passed to the analysis
+            routine. Set up differently per model type.
+        precision (np.ndarray): array of sampled measurement error precision values, populated in self.from_mcmc() after
+            the MCMC run is completed.
+
+    """
+
+    n_sensor: int = field(init=False)
+    precision_index: np.ndarray = field(init=False)
+    precision_parameter: parameter.Parameter = field(init=False)
+    prior_precision_shape: Union[np.ndarray, float] = field(init=False)
+    prior_precision_rate: Union[np.ndarray, float] = field(init=False)
+    initial_precision: Union[np.ndarray, float] = field(init=False)
+    precision: np.ndarray = field(init=False)
+
+    def initialise(
+        self, sensor_object: SensorGroup, meteorology: MeteorologyGroup = None, gas_species: GasSpecies = None
+    ):
+        """Take data inputs and extract relevant properties.
+
+        Args:
+            sensor_object (SensorGroup): sensor data.
+            meteorology (MeteorologyGroup): meteorology data. Defaults to None.
+            gas_species (GasSpecies): gas species information. Defaults to None.
+
+        """
+        self.n_sensor = sensor_object.nof_sensors
+
+    def make_model(self, model: list = None) -> list:
+        """Take model list and append new elements from current model component.
+
+        Args:
+            model (list, optional): Current list of model elements. Defaults to None.
+
+        Returns:
+            list: model output list.
+
+        """
+        if model is None:
+            model = []
+        model.append(Gamma("tau", shape="a_tau", rate="b_tau"))
+        return model
+
+    def make_sampler(self, model: Model, sampler_list: list = None) -> list:
+        """Take sampler list and append new elements from current model component.
+
+        Args:
+            model (Model): Full model list of distributions.
+            sampler_list (list, optional): Current list of samplers. Defaults to None.
+
+        Returns:
+            list: sampler output list.
+
+        """
+        if sampler_list is None:
+            sampler_list = []
+        sampler_list.append(NormalGamma("tau", model))
+        return sampler_list
+
+    def make_state(self, state: dict = None) -> dict:
+        """Take state dictionary and append initial values from model component.
+
+        Args:
+            state (dict, optional): current state vector. Defaults to None.
+
+        Returns:
+            dict: current state vector with components added.
+
+        """
+        if state is None:
+            state = {}
+        state["a_tau"] = self.prior_precision_shape.flatten()
+        state["b_tau"] = self.prior_precision_rate.flatten()
+        state["precision_index"] = self.precision_index
+        state["tau"] = self.initial_precision.flatten()
+        return state
+
+    def from_mcmc(self, store: dict):
+        """Extract results of mcmc from mcmc.store and attach to components.
+
+        Args:
+            store (dict): mcmc result dictionary.
+
+        """
+        self.precision = store["tau"]
+
+
+@dataclass
+class BySensor(ErrorModel):
+    """Version of measurement precision where each sensor object has a different precision.
+
+    Attributes:
+        prior_precision_shape (Union[np.ndarray, float]): prior shape parameters for the precision model, can be
+            specified either as a float or as a (nof_sensors, ) np.ndarray: a float specification will result in
+            the same parameter value for each sensor. Defaults to 1e-3.
+        prior_precision_rate (Union[np.ndarray, float]): prior rate parameters for the precision model, can be
+            specified either as a float or as a (nof_sensors, ) np.ndarray: a float specification will result in
+            the same parameter value for each sensor. Defaults to 1e-3.
+        initial_precision (Union[np.ndarray, float]): initial value for the precision parameters, can be specified
+            either as a float or as a (nof_sensors, ) np.ndarray: a float specification will result in the same
+            parameter value for each sensor. Defaults to 1.
+        precision_index (np.ndarray): index mapping precision parameters onto observations. Parameters 1:n_sensor are
+            mapped as the measurement error precisions of the corresponding sensors.
+        precision_parameter (Parameter.MixtureParameterMatrix): parameter specification for this model, maps the
+            current value of the parameter in the state dict onto the concentration data precisions.
+
+    """
+
+    prior_precision_shape: Union[np.ndarray, float] = 1e-3
+    prior_precision_rate: Union[np.ndarray, float] = 1e-3
+    initial_precision: Union[np.ndarray, float] = 1.0
+
+    def initialise(
+        self, sensor_object: SensorGroup, meteorology: MeteorologyGroup = None, gas_species: GasSpecies = None
+    ):
+        """Set up the error model using sensor properties.
+
+        Args:
+            sensor_object (SensorGroup): sensor data.
+            meteorology (MeteorologyGroup): meteorology data. Defaults to None.
+            gas_species (GasSpecies): gas species information. Defaults to None.
+
+        """
+        super().initialise(sensor_object=sensor_object, meteorology=meteorology, gas_species=gas_species)
+        self.prior_precision_shape = self.prior_precision_shape * np.ones((self.n_sensor,))
+        self.prior_precision_rate = self.prior_precision_rate * np.ones((self.n_sensor,))
+        self.initial_precision = self.initial_precision * np.ones((self.n_sensor,))
+        self.precision_index = sensor_object.sensor_index
+        self.precision_parameter = parameter.MixtureParameterMatrix(param="tau", allocation="precision_index")
+
+    def plot_iterations(self, plot: "Plot", sensor_object: Union[SensorGroup, Sensor], burn_in_value: int) -> "Plot":
+        """Plots the error model values for every sensor with respect to the MCMC iterations.
+
+        Args:
+            sensor_object (Union[SensorGroup, Sensor]): Sensor object associated with the error_model
+            burn_in_value (int): Burn in value to show in plot.
+            plot (Plot): Plot object to which this figure will be added in the figure dictionary
+
+        Returns:
+            plot (Plot): Plot object to which this figure is added in the figure dictionary with
+                key 'error_model_iterations'
+
+        """
+        plot.plot_trace_per_sensor(
+            object_to_plot=self, sensor_object=sensor_object, plot_type="line", burn_in=burn_in_value
+        )
+
+        return plot
+
+    def plot_distributions(self, plot: "Plot", sensor_object: Union[SensorGroup, Sensor], burn_in_value: int) -> "Plot":
+        """Plots the distribution of the error model values after the burn in for every sensor.
+
+        Args:
+            sensor_object (Union[SensorGroup, Sensor]): Sensor object associated with the error_model
+            burn_in_value (int): Burn in value to show in plot.
+            plot (Plot): Plot object to which this figure will be added in the figure dictionary
+
+        Returns:
+            plot (Plot): Plot object to which this figure is added in the figure dictionary with
+                key 'error_model_distributions'
+
+        """
+        plot.plot_trace_per_sensor(
+            object_to_plot=self, sensor_object=sensor_object, plot_type="box", burn_in=burn_in_value
+        )
+
+        return plot
+
+
+@dataclass
+class ByRelease(ErrorModel):
+    """ByRelease error model, special case of the measurement precision model.
+
+    Version of the measurement precision model where each sensor object has a different precision, and there are
+    different precisions for periods inside and outside controlled release periods. For all parameters: the first
+    element corresponds to the case where the sources are OFF; the second element corresponds to the case where the
+    sources are ON.
+
+    Attributes:
+        prior_precision_shape (np.ndarray): prior shape parameters for the precision model, can be
+            specified either as a (2, 1) np.ndarray or as a (2, nof_sensors) np.ndarray: the former specification
+            will result in the same prior specification for the off/on precisions for each sensor. Defaults to
+            np.array([1e-3, 1e-3]).
+        prior_precision_rate (np.ndarray): prior rate parameters for the precision model, can be
+            specified either as a (2, 1) np.ndarray or as a (2, nof_sensors) np.ndarray: the former specification
+            will result in the same prior specification for the off/on precisions for each sensor. Defaults to
+            np.array([1e-3, 1e-3]).
+        initial_precision (np.ndarray): initial value for the precision parameters, can be
+            specified either as a (2, 1) np.ndarray or as a (2, nof_sensors) np.ndarray: the former specification
+            will result in the same prior specification for the off/on precisions for each sensor. Defaults to
+            np.array([1.0, 1.0]).
+        precision_index (np.ndarray): index mapping precision parameters onto observations. Parameters 1:n_sensor are
+            mapped onto each sensor for the periods where the sources are OFF; parameters (n_sensor + 1):(2 * n_sensor)
+            are mapped onto each sensor for the periods where the sources are ON.
+        precision_parameter (Parameter.MixtureParameterMatrix): parameter specification for this model, maps the
+            current value of the parameter in the state dict onto the concentration data precisions.
+
+    """
+
+    prior_precision_shape: np.ndarray = field(default_factory=lambda: np.array([1e-3, 1e-3], ndmin=2).T)
+    prior_precision_rate: np.ndarray = field(default_factory=lambda: np.array([1e-3, 1e-3], ndmin=2).T)
+    initial_precision: np.ndarray = field(default_factory=lambda: np.array([1.0, 1.0], ndmin=2).T)
+
+    def initialise(
+        self, sensor_object: SensorGroup, meteorology: MeteorologyGroup = None, gas_species: GasSpecies = None
+    ):
+        """Set up the error model using sensor properties.
+
+        Args:
+            sensor_object (SensorGroup): sensor data.
+            meteorology (MeteorologyGroup): meteorology data. Defaults to None.
+            gas_species (GasSpecies): gas species information. Defaults to None.
+
+        """
+        super().initialise(sensor_object=sensor_object, meteorology=meteorology, gas_species=gas_species)
+        self.prior_precision_shape = self.prior_precision_shape * np.ones((2, self.n_sensor))
+        self.prior_precision_rate = self.prior_precision_rate * np.ones((2, self.n_sensor))
+        self.initial_precision = self.initial_precision * np.ones((2, self.n_sensor))
+        self.precision_index = sensor_object.sensor_index + sensor_object.source_on * self.n_sensor
+        self.precision_parameter = parameter.MixtureParameterMatrix(param="tau", allocation="precision_index")
+
+    def plot_iterations(self, plot: "Plot", sensor_object: Union[SensorGroup, Sensor], burn_in_value: int) -> "Plot":
+        """Plot the estimated error model parameters against iterations of the MCMC chain.
+
+        Works by simply creating a separate plot for each of the two categories of precision parameter (when the
+        sources are on/off). Creates a BySensor() object for each of the off/on precision cases, and then makes a
+        call to its plot function.
+
+        Args:
+            sensor_object (Union[SensorGroup, Sensor]): Sensor object associated with the error_model
+            burn_in_value (int): Burn in value to show in plot.
+            plot (Plot): Plot object to which this figure will be added in the figure dictionary
+
+        Returns:
+            plot (Plot): Plot object to which this figure is added in the figure dictionary with
+                key 'error_model_iterations'
+
+        """
+        figure_keys = ["error_model_off_iterations", "error_model_on_iterations"]
+        figure_titles = [
+            "Estimated error parameter values: sources off",
+            "Estimated error parameter values: sources on",
+        ]
+        precision_arrays = [
+            self.precision[: sensor_object.nof_sensors, :],
+            self.precision[sensor_object.nof_sensors :, :],
+        ]
+        for key, title, array in zip(figure_keys, figure_titles, precision_arrays):
+            error_model = BySensor()
+            error_model.precision = array
+            plot = error_model.plot_iterations(plot, sensor_object, burn_in_value)
+            plot.figure_dict[key] = plot.figure_dict.pop("error_model_iterations")
+            plot.figure_dict[key].update_layout(title=title)
+        return plot
+
+    def plot_distributions(self, plot: "Plot", sensor_object: Union[SensorGroup, Sensor], burn_in_value: int) -> "Plot":
+        """Plot the estimated distributions of error model parameters.
+
+        Works by simply creating a separate plot for each of the two categories of precision parameter (when the
+        sources are off/on). Creates a BySensor() object for each of the off/on precision cases, and then makes a
+        call to its plot function.
+
+        Args:
+            sensor_object (Union[SensorGroup, Sensor]): Sensor object associated with the error_model
+            burn_in_value (int): Burn in value to show in plot.
+            plot (Plot): Plot object to which this figure will be added in the figure dictionary
+
+        Returns:
+            plot (Plot): Plot object to which this figure is added in the figure dictionary with
+                key 'error_model_distributions'
+
+        """
+        figure_keys = ["error_model_off_distributions", "error_model_on_distributions"]
+        figure_titles = [
+            "Estimated error parameter distribution: sources off",
+            "Estimated error parameter distribution: sources on",
+        ]
+        precision_arrays = [
+            self.precision[: sensor_object.nof_sensors, :],
+            self.precision[sensor_object.nof_sensors :, :],
+        ]
+        for key, title, array in zip(figure_keys, figure_titles, precision_arrays):
+            error_model = BySensor()
+            error_model.precision = array
+            plot = error_model.plot_distributions(plot, sensor_object, burn_in_value)
+            plot.figure_dict[key] = plot.figure_dict.pop("error_model_distributions")
+            plot.figure_dict[key].update_layout(title=title)
+        return plot

pyelq/component/offset.py

@@ -0,0 +1,183 @@
+# SPDX-FileCopyrightText: 2024 Shell Global Solutions International B.V. All Rights Reserved.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+# -*- coding: utf-8 -*-
+"""Offset module."""
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Union
+
+import numpy as np
+from openmcmc import parameter
+from openmcmc.distribution.distribution import Gamma
+from openmcmc.distribution.location_scale import Normal
+from openmcmc.model import Model
+from openmcmc.sampler.sampler import NormalGamma, NormalNormal
+from scipy import sparse
+
+from pyelq.component.component import Component
+from pyelq.gas_species import GasSpecies
+from pyelq.meteorology import Meteorology
+from pyelq.sensor.sensor import Sensor, SensorGroup
+
+if TYPE_CHECKING:
+    from pyelq.plotting.plot import Plot
+
+
+@dataclass
+class PerSensor(Component):
+    """Offset implementation which assumes an additive offset between sensors.
+
+    The offset is which is constant in space and time and accounts for calibration differences between sensors.
+    To maintain parameter identifiability, the offset for the first sensor (with index 0) is assumed to be 0, and other
+    sensor offsets are defined relative to this beam.
+
+    Attributes:
+        n_sensor (int): number of sensors in the sensor object used for analysis.
+        offset (np.ndarray): array of sampled offset values, populated in self.from_mcmc() after the MCMC run is
+            completed.
+        precision_scalar (np.ndarray): array of sampled offset precision values, populated in self.from_mcmc() after
+            the MCMC run is completed. Only populated if update_precision is True.
+        indicator_basis (sparse.csc_matrix): [nof_observations x (nof_sensors - 1)] sparse matrix which assigns the
+            offset parameters to the correct observations.
+        update_precision (bool): logical indicating whether the offset prior precision parameter should be updated as
+            part of the analysis.
+        mean_offset (float): prior mean parameter for the offsets, assumed to be the same for each beam. Default is 0.
+        prior_precision_shape (float): shape parameter for the prior gamma distribution for the scalar precision
+            parameter. Default is 1e-3.
+        prior_precision_rate (float): rate parameter for the prior gamma distribution for the scalar precision
+            parameter(s). Default is 1e-3.
+        initial_precision (float): initial value for the scalar precision parameter. Default is 1.0.
+
+    """
+
+    n_sensor: int = field(init=False)
+    offset: np.ndarray = field(init=False)
+    precision_scalar: np.ndarray = field(init=False)
+    indicator_basis: sparse.csc_matrix = field(init=False)
+    update_precision: bool = False
+    mean_offset: float = 0.0
+    prior_precision_shape: float = 1e-3
+    prior_precision_rate: float = 1e-3
+    initial_precision: float = 1.0
+
+    def initialise(self, sensor_object: SensorGroup, meteorology: Meteorology, gas_species: GasSpecies):
+        """Take data inputs and extract relevant properties.
+
+        Args:
+            sensor_object (SensorGroup): sensor data
+            meteorology (MeteorologyGroup): meteorology data wind data
+            gas_species (GasSpecies): gas species information
+
+        """
+        self.n_sensor = len(sensor_object)
+        self.indicator_basis = sparse.csc_matrix(
+            np.equal(sensor_object.sensor_index[:, np.newaxis], np.array(range(1, self.n_sensor)))
+        )
+
+    def make_model(self, model: list = None) -> list:
+        """Take model list and append new elements from current model component.
+
+        Args:
+            model (list, optional): Current list of model elements. Defaults to [].
+
+        Returns:
+            list: model output list.
+
+        """
+        if model is None:
+            model = []
+        off_precision_predictor = parameter.ScaledMatrix(matrix="P_d", scalar="lambda_d")
+        model.append(Normal("d", mean="mu_d", precision=off_precision_predictor))
+        if self.update_precision:
+            model.append(Gamma("lambda_d", shape="a_lam_d", rate="b_lam_d"))
+        return model
+
+    def make_sampler(self, model: Model, sampler_list: list = None) -> list:
+        """Take sampler list and append new elements from current model component.
+
+        Args:
+            model (Model): Full model list of distributions.
+            sampler_list (list, optional): Current list of samplers. Defaults to [].
+
+        Returns:
+            list: sampler output list.
+
+        """
+        if sampler_list is None:
+            sampler_list = []
+        sampler_list.append(NormalNormal("d", model))
+        if self.update_precision:
+            sampler_list.append(NormalGamma("lambda_d", model))
+        return sampler_list
+
+    def make_state(self, state: dict = None) -> dict:
+        """Take state dictionary and append initial values from model component.
+
+        Args:
+            state (dict, optional): current state vector. Defaults to {}.
+
+        Returns:
+            dict: current state vector with components added.
+
+        """
+        if state is None:
+            state = {}
+        state["mu_d"] = np.ones((self.n_sensor - 1, 1)) * self.mean_offset
+        state["d"] = np.zeros((self.n_sensor - 1, 1))
+        state["B_d"] = self.indicator_basis
+        state["P_d"] = sparse.eye(self.n_sensor - 1, format="csc")
+        state["lambda_d"] = self.initial_precision
+        if self.update_precision:
+            state["a_lam_d"] = self.prior_precision_shape
+            state["b_lam_d"] = self.prior_precision_rate
+        return state
+
+    def from_mcmc(self, store: dict):
+        """Extract results of mcmc from mcmc.store and attach to components.
+
+        Args:
+            store (dict): mcmc result dictionary.
+
+        """
+        self.offset = store["d"]
+        if self.update_precision:
+            self.precision_scalar = store["lambda_d"]
+
+    def plot_iterations(self, plot: "Plot", sensor_object: Union[SensorGroup, Sensor], burn_in_value: int) -> "Plot":
+        """Plots the offset values for every sensor with respect to the MCMC iterations.
+
+        Args:
+            sensor_object (Union[SensorGroup, Sensor]): Sensor object associated with the offset_model
+            burn_in_value (int): Burn in value to show in plot.
+            plot (Plot): Plot object to which this figure will be added in the figure dictionary
+
+        Returns:
+            plot (Plot): Plot object to which this figure is added in the figure dictionary with
+                key 'offset_iterations'
+
+        """
+        plot.plot_trace_per_sensor(
+            object_to_plot=self, sensor_object=sensor_object, plot_type="line", burn_in=burn_in_value
+        )
+
+        return plot
+
+    def plot_distributions(self, plot: "Plot", sensor_object: Union[SensorGroup, Sensor], burn_in_value: int) -> "Plot":
+        """Plots the distribution of the offset values after the burn in for every sensor.
+
+        Args:
+            sensor_object (Union[SensorGroup, Sensor]): Sensor object associated with the offset_model
+            burn_in_value (int): Burn in value to use for plot.
+            plot (Plot): Plot object to which this figure will be added in the figure dictionary
+
+        Returns:
+            plot (Plot): Plot object to which this figure is added in the figure dictionary with
+                key 'offset_distributions'
+
+        """
+        plot.plot_trace_per_sensor(
+            object_to_plot=self, sensor_object=sensor_object, plot_type="box", burn_in=burn_in_value
+        )
+
+        return plot