pyelq 1.1.0__py3-none-any.whl

pyelq/component/source_model.py

@@ -0,0 +1,875 @@
+# SPDX-FileCopyrightText: 2024 Shell Global Solutions International B.V. All Rights Reserved.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+# -*- coding: utf-8 -*-
+"""Component Class and subclasses for source model.
+
+A SourceModel instance inherits from 3 super-classes:
+    - Component: this is the general superclass for ELQModel components, which prototypes generic methods.
+    - A type of SourceGrouping: this class type implements an allocation of sources to different categories (e.g. slab
+        or spike), and sets up a sampler for estimating the classification of each source within the source map.
+        Inheriting from the NullGrouping class ensures that the allocation of all sources is fixed during the inversion,
+        and is not updated.
+    - A type of SourceDistribution: this class type implements a particular type of response distribution (mostly
+        Normal, but also allows for cases where we have e.g. exp(log_s) or a non-Gaussian prior).
+
+"""
+
+from abc import abstractmethod
+from copy import deepcopy
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Tuple, Union
+
+import numpy as np
+from openmcmc import parameter
+from openmcmc.distribution.distribution import Categorical, Gamma, Poisson, Uniform
+from openmcmc.distribution.location_scale import Normal as mcmcNormal
+from openmcmc.model import Model
+from openmcmc.sampler.metropolis_hastings import RandomWalkLoop
+from openmcmc.sampler.reversible_jump import ReversibleJump
+from openmcmc.sampler.sampler import MixtureAllocation, NormalGamma, NormalNormal
+
+from pyelq.component.component import Component
+from pyelq.coordinate_system import ENU
+from pyelq.dispersion_model.gaussian_plume import GaussianPlume
+from pyelq.gas_species import GasSpecies
+from pyelq.meteorology import Meteorology
+from pyelq.sensor.sensor import SensorGroup
+from pyelq.source_map import SourceMap
+
+if TYPE_CHECKING:
+    from pyelq.plotting.plot import Plot
+
+
+@dataclass
+class SourceGrouping:
+    """Superclass for source grouping approach.
+
+    Source grouping method determines the group allocation of each source in the model, e.g: slab and spike
+    distribution makes an on/off allocation for each source.
+
+    Attributes:
+        nof_sources (int): number of sources in the model.
+        emission_rate_mean (Union[float, np.ndarray]): prior mean parameter for the emission rate distribution.
+        _source_key (str): label for the source parameter to be used in the distributions, samplers, MCMC state etc.
+
+    """
+
+    nof_sources: int = field(init=False)
+    emission_rate_mean: Union[float, np.ndarray] = field(init=False)
+    _source_key: str = field(init=False, default="s")
+
+    @abstractmethod
+    def make_allocation_model(self, model: list) -> list:
+        """Initialise the source allocation part of the model, and the parameters of the source response distribution.
+
+        Args:
+            model (list): overall model, consisting of list of distributions.
+
+        Returns:
+            list: overall model list, updated with allocation distribution.
+
+        """
+
+    @abstractmethod
+    def make_allocation_sampler(self, model: Model, sampler_list: list) -> list:
+        """Initialise the allocation part of the sampler.
+
+        Args:
+            model (Model): overall model, consisting of list of distributions.
+            sampler_list (list): list of samplers for individual parameters.
+
+        Returns:
+            list: sampler_list updated with sampler for the source allocation.
+
+        """
+
+    @abstractmethod
+    def make_allocation_state(self, state: dict) -> dict:
+        """Initialise the allocation part of the state.
+
+        Args:
+            state (dict): dictionary containing current state information.
+
+        Returns:
+            dict: state updated with parameters related to the source grouping.
+
+        """
+
+    @abstractmethod
+    def from_mcmc_group(self, store: dict):
+        """Extract posterior allocation samples from the MCMC sampler, attach them to the class.
+
+        Args:
+            store (dict): dictionary containing samples from the MCMC.
+
+        """
+
+
+@dataclass
+class NullGrouping(SourceGrouping):
+    """Null grouping: the grouping of the sources will not change during the course of the inversion.
+
+    Note that this is intended to support two distinct cases:
+        1) The case where the source map is fixed, and a given prior mean and prior precision value are assigned to
+            each source (can be a common value for all sources, or can be a distinct allocation to each element of the
+            source map).
+        2) The case where the dimensionality of the source map is changing during the inversion, and a common prior
+            mean and precision term are used for all sources.
+
+    """
+
+    def make_allocation_model(self, model: list) -> list:
+        """Initialise the source allocation part of the model.
+
+        In the NullGrouping case, the source allocation is fixed throughout, so this function does nothing (simply
+        returns the existing model un-modified).
+
+        Args:
+            model (list): model as constructed so far, consisting of list of distributions.
+
+        Returns:
+            list: overall model list, updated with allocation distribution.
+
+        """
+        return model
+
+    def make_allocation_sampler(self, model: Model, sampler_list: list) -> list:
+        """Initialise the allocation part of the sampler.
+
+        In the NullGrouping case, the source allocation is fixed throughout, so this function does nothing (simply
+        returns the existing sampler list un-modified).
+
+        Args:
+            model (Model): overall model set for the problem.
+            sampler_list (list): list of samplers for individual parameters.
+
+        Returns:
+            list: sampler_list updated with sampler for the source allocation.
+
+        """
+        return sampler_list
+
+    def make_allocation_state(self, state: dict) -> dict:
+        """Initialise the allocation part of the state.
+
+        The prior mean parameter and the fixed source allocation are added to the state.
+
+        Args:
+            state (dict): dictionary containing current state information.
+
+        Returns:
+            dict: state updated with parameters related to the source grouping.
+
+        """
+        state["mu_s"] = np.array(self.emission_rate_mean, ndmin=1)
+        state["alloc_s"] = np.zeros((self.nof_sources, 1), dtype="int")
+        return state
+
+    def from_mcmc_group(self, store: dict):
+        """Extract posterior allocation samples from the MCMC sampler, attach them to the class.
+
+        We have not implemented anything here as there is nothing to fetch from the MCMC solution here for the
+        NullGrouping Class.
+
+        Args:
+            store (dict): dictionary containing samples from the MCMC.
+
+        """
+
+
+@dataclass
+class SlabAndSpike(SourceGrouping):
+    """Slab and spike source model, special case for the source grouping.
+
+    Slab and spike: the prior for the emission rates is a two-component mixture, and the allocation is to be
+    estimated as part of the inversion.
+
+    Attributes:
+        slab_probability (float): prior probability of allocation to the slab component. Defaults to 0.05.
+        allocation (np.ndarray): set of allocation samples, with shape=(n_sources, n_iterations). Attached to
+            the class by self.from_mcmc_group().
+
+    """
+
+    slab_probability: float = 0.05
+    allocation: np.ndarray = field(init=False)
+
+    def make_allocation_model(self, model: list) -> list:
+        """Initialise the source allocation part of the model.
+
+        Args:
+            model (list): model as constructed so far, consisting of list of distributions.
+
+        Returns:
+            list: overall model list, updated with allocation distribution.
+
+        """
+        model.append(Categorical("alloc_s", prob="s_prob"))
+        return model
+
+    def make_allocation_sampler(self, model: Model, sampler_list: list) -> list:
+        """Initialise the allocation part of the sampler.
+
+        Args:
+            model (Model): overall model set for the problem.
+            sampler_list (list): list of samplers for individual parameters.
+
+        Returns:
+            list: sampler_list updated with sampler for the source allocation.
+
+        """
+        sampler_list.append(MixtureAllocation(param="alloc_s", model=model, response_param=self._source_key))
+        return sampler_list
+
+    def make_allocation_state(self, state: dict) -> dict:
+        """Initialise the allocation part of the state.
+
+        Args:
+            state (dict): dictionary containing current state information.
+
+        Returns:
+            dict: state updated with parameters related to the source grouping.
+
+        """
+        state["mu_s"] = np.array(self.emission_rate_mean, ndmin=1)
+        state["s_prob"] = np.tile(np.array([self.slab_probability, 1 - self.slab_probability]), (self.nof_sources, 1))
+        state["alloc_s"] = np.ones((self.nof_sources, 1), dtype="int")
+        return state
+
+    def from_mcmc_group(self, store: dict):
+        """Extract posterior allocation samples from the MCMC sampler, attach them to the class.
+
+        Args:
+            store (dict): dictionary containing samples from the MCMC.
+
+        """
+        self.allocation = store["alloc_s"]
+
+
+@dataclass
+class SourceDistribution:
+    """Superclass for source emission rate distribution.
+
+    Source distribution determines the type of prior to be used for the source emission rates, and the transformation
+    linking the source parameters and the data.
+
+    Elements related to transformation of source parameters are also specified at the model level.
+
+    Attributes:
+        nof_sources (int): number of sources in the model.
+        emission_rate (np.ndarray): set of emission rate samples, with shape=(n_sources, n_iterations). Attached to
+            the class by self.from_mcmc_dist().
+
+    """
+
+    nof_sources: int = field(init=False)
+    emission_rate: np.ndarray = field(init=False)
+
+    @abstractmethod
+    def make_source_model(self, model: list) -> list:
+        """Add distributional component to the overall model corresponding to the source emission rate distribution.
+
+        Args:
+            model (list): model as constructed so far, consisting of list of distributions.
+
+        Returns:
+            list: overall model list, updated with distributions related to source prior.
+
+        """
+
+    @abstractmethod
+    def make_source_sampler(self, model: Model, sampler_list: list) -> list:
+        """Initialise the source prior distribution part of the sampler.
+
+        Args:
+            model (Model): overall model set for the problem.
+            sampler_list (list): list of samplers for individual parameters.
+
+        Returns:
+            list: sampler_list updated with sampler for the emission rate parameters.
+
+        """
+
+    @abstractmethod
+    def make_source_state(self, state: dict) -> dict:
+        """Initialise the emission rate parts of the state.
+
+        Args:
+            state (dict): dictionary containing current state information.
+
+        Returns:
+            dict: state updated with parameters related to the source emission rates.
+
+        """
+
+    @abstractmethod
+    def from_mcmc_dist(self, store: dict):
+        """Extract posterior emission rate samples from the MCMC, attach them to the class.
+
+        Args:
+            store (dict): dictionary containing samples from the MCMC.
+
+        """
+
+
+@dataclass
+class NormalResponse(SourceDistribution):
+    """(Truncated) Gaussian prior for sources.
+
+    No transformation applied to parameters, i.e.:
+    - Prior distribution: s ~ N(mu, 1/precision)
+    - Likelihood contribution: y = A*s + b + ...
+
+    Attributes:
+        truncation (bool): indication of whether the emission rate prior should be truncated at 0. Defaults to True.
+        emission_rate_lb (Union[float, np.ndarray]): lower bound for the source emission rates. Defaults to 0.
+        emission_rate_mean (Union[float, np.ndarray]): prior mean for the emission rate distribution. Defaults to 0.
+
+    """
+
+    truncation: bool = True
+    emission_rate_lb: Union[float, np.ndarray] = 0
+    emission_rate_mean: Union[float, np.ndarray] = 0
+
+    def make_source_model(self, model: list) -> list:
+        """Add distributional component to the overall model corresponding to the source emission rate distribution.
+
+        Args:
+            model (list): model as constructed so far, consisting of list of distributions.
+
+        Returns:
+            list: model, updated with distributions related to source prior.
+
+        """
+        domain_response_lower = None
+        if self.truncation:
+            domain_response_lower = self.emission_rate_lb
+
+        model.append(
+            mcmcNormal(
+                "s",
+                mean=parameter.MixtureParameterVector(param="mu_s", allocation="alloc_s"),
+                precision=parameter.MixtureParameterMatrix(param="lambda_s", allocation="alloc_s"),
+                domain_response_lower=domain_response_lower,
+            )
+        )
+        return model
+
+    def make_source_sampler(self, model: Model, sampler_list: list = None) -> list:
+        """Initialise the source prior distribution part of the sampler.
+
+        Args:
+            model (Model): overall model set for the problem.
+            sampler_list (list): list of samplers for individual parameters.
+
+        Returns:
+            list: sampler_list updated with sampler for the emission rate parameters.
+
+        """
+        if sampler_list is None:
+            sampler_list = []
+        sampler_list.append(NormalNormal("s", model))
+        return sampler_list
+
+    def make_source_state(self, state: dict) -> dict:
+        """Initialise the emission rate part of the state.
+
+        Args:
+            state (dict): dictionary containing current state information.
+
+        Returns:
+            dict: state updated with initial emission rate vector.
+
+        """
+        state["s"] = np.zeros((self.nof_sources, 1))
+        return state
+
+    def from_mcmc_dist(self, store: dict):
+        """Extract posterior emission rate samples from the MCMC sampler, attach them to the class.
+
+        Args:
+            store (dict): dictionary containing samples from the MCMC.
+
+        """
+        self.emission_rate = store["s"]
+
+
+@dataclass
+class SourceModel(Component, SourceGrouping, SourceDistribution):
+    """Superclass for the specification of the source model in an inversion run.
+
+    Various different types of model. A SourceModel is an optional component of a model, and thus inherits
+    from Component.
+
+    A subclass instance of SourceModel must inherit from:
+        - an INSTANCE of SourceDistribution, which specifies a prior emission rate distribution for all sources in the
+            source map.
+        - an INSTANCE of SourceGrouping, which specifies a type of mixture prior specification for the sources (for
+            which the allocation is to be estimated as part of the inversion).
+
+    If the flag reversible_jump == True, then the number of sources and their locations are also estimated as part of
+    the inversion, in addition to the emission rates. If this flag is set to true, the sensor_object, meteorology and
+    gas_species objects are all attached to the class, as they will be required in the repeated computation of updates
+    to the coupling matrix during the inversion.
+
+    Attributes:
+        dispersion_model (GaussianPlume): dispersion model used to generate the couplings between source locations and
+            sensor observations.
+        coupling (np.ndarray): coupling matrix generated using dispersion_model.
+
+        sensor_object (SensorGroup): stores sensor information for reversible jump coupling updates.
+        meteorology (MeteorologyGroup): stores meteorology information for reversible jump coupling updates.
+        gas_species (GasSpecies): stores gas species information for reversible jump coupling updates.
+
+        reversible_jump (bool): logical indicating whether the reversible jump algorithm for estimation of the number
+            of sources and their locations should be run. Defaults to False.
+        random_walk_step_size (np.ndarray): (3 x 1) array specifying the standard deviations of the distributions
+            from which the random walk sampler draws new source locations. Defaults to np.array([1.0, 1.0, 0.1]).
+        site_limits (np.ndarray): (3 x 2) array specifying the lower (column 0) and upper (column 1) limits of the
+            analysis site. Only relevant for cases where reversible_jump == True (where sources are free to move in
+            the solution).
+        rate_num_sources (int): specification for the parameter for the Poisson prior distribution for the total number
+            of sources. Only relevant for cases where reversible_jump == True (where the number of sources in the
+            solution can change).
+        n_sources_max (int): maximum number of sources that can feature in the solution. Only relevant for cases where
+            reversible_jump == True (where the number of sources in the solution can change).
+        emission_proposal_std (float): standard deviation of the truncated Gaussian distribution used to propose the
+            new source emission rate in case of a birth move.
+
+        update_precision (bool): logical indicating whether the prior precision parameter for emission rates should be
+            updated as part of the inversion. Defaults to false.
+        prior_precision_shape (Union[float, np.ndarray]): shape parameters for the prior Gamma distribution for the
+            source precision parameter.
+        prior_precision_rate (Union[float, np.ndarray]): rate parameters for the prior Gamma distribution for the
+            source precision parameter.
+        initial_precision (Union[float, np.ndarray]): initial value for the source emission rate precision parameter.
+        precision_scalar (np.ndarray): precision values generated by MCMC inversion.
+
+        coverage_detection (float): sensor detection threshold (in ppm) to be used for coverage calculations.
+        coverage_test_source (float): test source (in kg/hr) which we wish to be able to see in coverage calculation.
+
+        threshold_function (Callable): Callable function which returns a single value that defines the threshold
+            for the coupling in a lambda function form. Examples: lambda x: np.quantile(x, 0.95, axis=0),
+            lambda x: np.max(x, axis=0), lambda x: np.mean(x, axis=0). Defaults to np.quantile.
+
+    """
+
+    dispersion_model: GaussianPlume = field(init=False, default=None)
+    coupling: np.ndarray = field(init=False)
+
+    sensor_object: SensorGroup = field(init=False, default=None)
+    meteorology: Meteorology = field(init=False, default=None)
+    gas_species: GasSpecies = field(init=False, default=None)
+
+    reversible_jump: bool = False
+    random_walk_step_size: np.ndarray = field(default_factory=lambda: np.array([1.0, 1.0, 0.1], ndmin=2).T)
+    site_limits: np.ndarray = None
+    rate_num_sources: int = 5
+    n_sources_max: int = 20
+    emission_proposal_std: float = 0.5
+
+    update_precision: bool = False
+    prior_precision_shape: Union[float, np.ndarray] = 1e-3
+    prior_precision_rate: Union[float, np.ndarray] = 1e-3
+    initial_precision: Union[float, np.ndarray] = 1.0
+    precision_scalar: np.ndarray = field(init=False)
+
+    coverage_detection: float = 0.1
+    coverage_test_source: float = 6.0
+
+    threshold_function: callable = lambda x: np.quantile(x, 0.95, axis=0)
+
+    @property
+    def nof_sources(self):
+        """Get number of sources in the source map."""
+        return self.dispersion_model.source_map.nof_sources
+
+    @property
+    def coverage_threshold(self):
+        """Compute coverage threshold from detection threshold and test source strength."""
+        return self.coverage_test_source / self.coverage_detection
+
+    def initialise(self, sensor_object: SensorGroup, meteorology: Meteorology, gas_species: GasSpecies):
+        """Set up the source model.
+
+        Extract required information from the sensor, meteorology and gas species objects:
+            - Attach coupling calculated using self.dispersion_model.
+            - (If self.reversible_jump == True) Attach objects to source model which will be used in RJMCMC sampler,
+                they will be required when we need to update the couplings when new source locations are proposed when
+                we move/birth/death.
+
+        Args:
+            sensor_object (SensorGroup): object containing sensor data.
+            meteorology (MeteorologyGroup): object containing meteorology data.
+            gas_species (GasSpecies): object containing gas species information.
+
+        """
+        self.initialise_dispersion_model(sensor_object)
+        self.coupling = self.dispersion_model.compute_coupling(
+            sensor_object, meteorology, gas_species, output_stacked=True
+        )
+        self.screen_coverage()
+        if self.reversible_jump:
+            self.sensor_object = sensor_object
+            self.meteorology = meteorology
+            self.gas_species = gas_species
+
+    def initialise_dispersion_model(self, sensor_object: SensorGroup):
+        """Initialise the dispersion model.
+
+        If a dispersion_model has already been attached to this instance, then this function takes no action.
+
+        If a dispersion_model has not already been attached to the instance, then this function adds a GaussianPlume
+        dispersion model, with a default source map that has limits set based on the sensor locations.
+
+        Args:
+            sensor_object (SensorGroup): object containing sensor data.
+
+        """
+        if self.dispersion_model is None:
+            source_map = SourceMap()
+            sensor_locations = sensor_object.location.to_enu()
+            location_object = ENU(
+                ref_latitude=sensor_locations.ref_latitude,
+                ref_longitude=sensor_locations.ref_longitude,
+                ref_altitude=sensor_locations.ref_altitude,
+            )
+            source_map.generate_sources(
+                coordinate_object=location_object,
+                sourcemap_limits=np.array(
+                    [
+                        [np.min(sensor_locations.east), np.max(sensor_locations.east)],
+                        [np.min(sensor_locations.north), np.max(sensor_locations.north)],
+                        [np.min(sensor_locations.up), np.max(sensor_locations.up)],
+                    ]
+                ),
+                sourcemap_type="grid",
+            )
+            self.dispersion_model = GaussianPlume(source_map)
+
+    def screen_coverage(self):
+        """Screen the initial source map for coverage."""
+        in_coverage_area = self.dispersion_model.compute_coverage(
+            self.coupling, coverage_threshold=self.coverage_threshold, threshold_function=self.threshold_function
+        )
+        self.coupling = self.coupling[:, in_coverage_area]
+        all_locations = self.dispersion_model.source_map.location.to_array()
+        screened_locations = all_locations[in_coverage_area, :]
+        self.dispersion_model.source_map.location.from_array(screened_locations)
+
+    def update_coupling_column(self, state: dict, update_column: int) -> dict:
+        """Update the coupling, based on changes to the source locations as part of inversion.
+
+        To be used in two different situations:
+            - movement of source locations (e.g. Metropolis Hastings, random walk).
+            - adding of new source locations (e.g. reversible jump birth move).
+        If [update_column < A.shape[1]]: an existing column of the A matrix is updated.
+        If [update_column == A.shape[1]]: a new column is appended to the right-hand side of the A matrix
+        (corresponding to a new source).
+
+        A central assumption of this function is that the sensor information and meteorology information
+        have already been interpolated onto the same space/time points.
+
+        If an update_column is supplied, the coupling for that source location only is calculated to save on
+        computation time. If update_column is None, then we just re-compute the whole coupling matrix.
+
+        Args:
+            state (dict): dictionary containing state parameters.
+            update_column (int): index of the coupling column to be updated.
+
+        Returns:
+            state (dict): state dictionary containing updated coupling information.
+
+        """
+        self.dispersion_model.source_map.location.from_array(state["z_src"][:, [update_column]].T)
+        new_coupling = self.dispersion_model.compute_coupling(
+            self.sensor_object, self.meteorology, self.gas_species, output_stacked=True, run_interpolation=False
+        )
+
+        if update_column == state["A"].shape[1]:
+            state["A"] = np.concatenate((state["A"], new_coupling), axis=1)
+        elif update_column < state["A"].shape[1]:
+            state["A"][:, [update_column]] = new_coupling
+        else:
+            raise ValueError("Invalid column specification for updating.")
+        return state
+
+    def birth_function(self, current_state: dict, prop_state: dict) -> Tuple[dict, float, float]:
+        """Update MCMC state based on source birth proposal.
+
+        Proposed state updated as follows:
+            1- Add column to coupling matrix for new source location.
+            2- If required, adjust other components of the state which correspond to the sources.
+        The source emission rate vector will be adjusted using the standardised functionality
+        in the openMCMC package.
+
+        After the coupling has been updated, a coverage test is applied for the new source
+        location. If the max coupling is too small, a large contribution is added to the
+        log-proposal density for the new state, to force the sampler to reject it.
+
+        A central assumption of this function is that the sensor information and meteorology information
+        have already been interpolated onto the same space/time points.
+
+        This function assumes that the new source location has been added as the final column of
+        the source location matrix, and so will correspondingly append the new coupling column to the right
+        hand side of the current state coupling, and append an emission rate as the last element of the
+        current state emission rate vector.
+
+        Args:
+            current_state (dict): dictionary containing parameters of the current state.
+            prop_state (dict): dictionary containing the parameters of the proposed state.
+
+        Returns:
+            prop_state (dict): proposed state, with coupling matrix and source emission rate vector updated.
+            logp_pr_g_cr (float): log-transition density of the proposed state given the current state
+                (i.e. log[p(proposed | current)])
+            logp_cr_g_pr (float): log-transition density of the current state given the proposed state
+                (i.e. log[p(current | proposed)])
+
+        """
+        prop_state = self.update_coupling_column(prop_state, int(prop_state["n_src"]) - 1)
+        prop_state["alloc_s"] = np.concatenate((prop_state["alloc_s"], np.array([0], ndmin=2)), axis=0)
+        in_cov_area = self.dispersion_model.compute_coverage(
+            prop_state["A"][:, -1],
+            coverage_threshold=self.coverage_threshold,
+            threshold_function=self.threshold_function,
+        )
+        if not in_cov_area:
+            logp_pr_g_cr = 1e10
+        else:
+            logp_pr_g_cr = 0.0
+        logp_cr_g_pr = 0.0
+
+        return prop_state, logp_pr_g_cr, logp_cr_g_pr
+
+    @staticmethod
+    def death_function(current_state: dict, prop_state: dict, deletion_index: int) -> Tuple[dict, float, float]:
+        """Update MCMC state based on source death proposal.
+
+        Proposed state updated as follows:
+            1- Remove column from coupling for deleted source.
+            2- If required, adjust other components of the state which correspond to the sources.
+        The source emission rate vector will be adjusted using the standardised functionality in the general_mcmc repo.
+
+        A central assumption of this function is that the sensor information and meteorology information have already
+        been interpolated onto the same space/time points.
+
+        Args:
+            current_state (dict): dictionary containing parameters of the current state.
+            prop_state (dict): dictionary containing the parameters of the proposed state.
+            deletion_index (int): index of the source to be deleted in the overall set of sources.
+
+        Returns:
+            prop_state (dict): proposed state, with coupling matrix and source emission rate vector updated.
+            logp_pr_g_cr (float): log-transition density of the proposed state given the current state
+                (i.e. log[p(proposed | current)])
+            logp_cr_g_pr (float): log-transition density of the current state given the proposed state
+                (i.e. log[p(current | proposed)])
+
+        """
+        prop_state["A"] = np.delete(prop_state["A"], obj=deletion_index, axis=1)
+        prop_state["alloc_s"] = np.delete(prop_state["alloc_s"], obj=deletion_index, axis=0)
+        logp_pr_g_cr = 0.0
+        logp_cr_g_pr = 0.0
+
+        return prop_state, logp_pr_g_cr, logp_cr_g_pr
+
+    def move_function(self, current_state: dict, update_column: int) -> dict:
+        """Re-compute the coupling after a source location move.
+
+        Function first updates the coupling column, and then checks whether the location passes a coverage test. If the
+        location does not have good enough coverage, the state reverts to the coupling from the current state.
+
+        Args:
+            current_state (dict): dictionary containing parameters of the current state.
+            update_column (int): index of the coupling column to be updated.
+
+        Returns:
+            dict: proposed state, with updated coupling matrix.
+
+        """
+        prop_state = deepcopy(current_state)
+        prop_state = self.update_coupling_column(prop_state, update_column)
+        in_cov_area = self.dispersion_model.compute_coverage(
+            prop_state["A"][:, update_column],
+            coverage_threshold=self.coverage_threshold,
+            threshold_function=self.threshold_function,
+        )
+        if not in_cov_area:
+            prop_state = deepcopy(current_state)
+        return prop_state
+
+    def make_model(self, model: list) -> list:
+        """Take model list and append new elements from current model component.
+
+        Args:
+            model (list): Current list of model elements.
+
+        Returns:
+            list: model list updated with source-related distributions.
+
+        """
+        model = self.make_allocation_model(model)
+        model = self.make_source_model(model)
+        if self.update_precision:
+            model.append(Gamma("lambda_s", shape="a_lam_s", rate="b_lam_s"))
+        if self.reversible_jump:
+            model.append(
+                Uniform(
+                    response="z_src",
+                    domain_response_lower=self.site_limits[:, [0]],
+                    domain_response_upper=self.site_limits[:, [1]],
+                )
+            )
+            model.append(Poisson(response="n_src", rate="rho"))
+        return model
+
+    def make_sampler(self, model: Model, sampler_list: list) -> list:
+        """Take sampler list and append new elements from current model component.
+
+        Args:
+            model (Model): Full model list of distributions.
+            sampler_list (list): Current list of samplers.
+
+        Returns:
+            list: sampler list updated with source-related samplers.
+
+        """
+        sampler_list = self.make_source_sampler(model, sampler_list)
+        sampler_list = self.make_allocation_sampler(model, sampler_list)
+        if self.update_precision:
+            sampler_list.append(NormalGamma("lambda_s", model))
+        if self.reversible_jump:
+            sampler_list = self.make_sampler_rjmcmc(model, sampler_list)
+        return sampler_list
+
+    def make_state(self, state: dict) -> dict:
+        """Take state dictionary and append initial values from model component.
+
+        Args:
+            state (dict): current state vector.
+
+        Returns:
+            dict: current state vector with source-related parameters added.
+
+        """
+        state = self.make_allocation_state(state)
+        state = self.make_source_state(state)
+        state["A"] = self.coupling
+        state["lambda_s"] = np.array(self.initial_precision, ndmin=1)
+        if self.update_precision:
+            state["a_lam_s"] = np.ones_like(self.initial_precision) * self.prior_precision_shape
+            state["b_lam_s"] = np.ones_like(self.initial_precision) * self.prior_precision_rate
+        if self.reversible_jump:
+            state["z_src"] = self.dispersion_model.source_map.location.to_array().T
+            state["n_src"] = state["z_src"].shape[1]
+            state["rho"] = self.rate_num_sources
+        return state
+
+    def make_sampler_rjmcmc(self, model: Model, sampler_list: list) -> list:
+        """Create the parts of the sampler related to the reversible jump MCMC scheme.
+
+        RJ MCMC scheme:
+            - create the RandomWalkLoop sampler object which updates the source locations one-at-a-time.
+            - create the ReversibleJump sampler which proposes birth/death moves to add/remove sources from the source
+                map.
+
+        Args:
+            model (Model): model object containing probability density objects for all uncertain
+                parameters.
+            sampler_list (list): list of existing samplers.
+
+        Returns:
+            sampler_list (list): list of samplers updated with samplers corresponding to RJMCMC routine.
+
+        """
+        sampler_list[-1].max_variable_size = self.n_sources_max
+
+        sampler_list.append(
+            RandomWalkLoop(
+                "z_src",
+                model,
+                step=self.random_walk_step_size,
+                max_variable_size=(3, self.n_sources_max),
+                domain_limits=self.site_limits,
+                state_update_function=self.move_function,
+            )
+        )
+        matching_params = {"variable": "s", "matrix": "A", "scale": 1.0, "limits": [0.0, 1e6]}
+        sampler_list.append(
+            ReversibleJump(
+                "n_src",
+                model,
+                step=np.array([1.0], ndmin=2),
+                associated_params="z_src",
+                n_max=self.n_sources_max,
+                state_birth_function=self.birth_function,
+                state_death_function=self.death_function,
+                matching_params=matching_params,
+            )
+        )
+        return sampler_list
+
+    def from_mcmc(self, store: dict):
+        """Extract results of mcmc from mcmc.store and attach to components.
+
+        Args:
+            store (dict): mcmc result dictionary.
+
+        """
+        self.from_mcmc_group(store)
+        self.from_mcmc_dist(store)
+        if self.update_precision:
+            self.precision_scalar = store["lambda_s"]
+
+    def plot_iterations(self, plot: "Plot", burn_in_value: int, y_axis_type: str = "linear") -> "Plot":
+        """Plot the emission rate estimates source model object against MCMC iteration.
+
+        Args:
+            burn_in_value (int): Burn in value to show in plot.
+            y_axis_type (str, optional): String to indicate whether the y-axis should be linear of log scale.
+            plot (Plot): Plot object to which this figure will be added in the figure dictionary.
+
+        Returns:
+            plot (Plot): Plot object to which the figures added in the figure dictionary with
+                keys 'estimated_values_plot'/'log_estimated_values_plot' and 'number_of_sources_plot'
+
+        """
+        plot.plot_emission_rate_estimates(source_model_object=self, burn_in=burn_in_value, y_axis_type=y_axis_type)
+        plot.plot_single_trace(object_to_plot=self)
+        return plot
+
+
+@dataclass
+class Normal(SourceModel, NullGrouping, NormalResponse):
+    """Normal model, with null allocation.
+
+    (Truncated) Gaussian prior for emission rates, no grouping/allocation; no transformation applied to emission rate
+    parameters.
+
+    Can be used in the following cases:
+        - Fixed set of sources (grid or specific locations), all with the same Gaussian prior distribution.
+        - Variable number of sources, with a common prior distribution, estimated using reversible jump MCMC.
+        - Fixed set of sources with a bespoke prior per source (using the allocation to map prior parameters onto
+            sources).
+
+    """
+
+
+@dataclass
+class NormalSlabAndSpike(SourceModel, SlabAndSpike, NormalResponse):
+    """Normal Slab and Spike model.
+
+    (Truncated) Gaussian prior for emission rates, slab and spike prior, with allocation estimation; no transformation
+    applied to emission rate parameters.
+
+    Attributes:
+        initial_precision (np.ndarray): initial precision parameter for a slab and spike case. shape=(2, 1).
+        emission_rate_mean (np.ndarray): emission rate prior mean for a slab and spike case. shape=(2, 1).
+
+    """
+
+    initial_precision: np.ndarray = field(default_factory=lambda: np.array([1 / (10**2), 1 / (0.01**2)], ndmin=2).T)
+    emission_rate_mean: np.ndarray = field(default_factory=lambda: np.array([0, 0], ndmin=2).T)