Repository: cryotools/cosipy (language: python)

Issue:
Consistent object oriented code
Rewrite the code to consistent OOP.


Discussion / hints:

Fix patch:
diff --git a/.coveragerc b/.coveragerc
index 1a1f830f..ebba4cc8 100644
--- a/.coveragerc
+++ b/.coveragerc
@@ -1,9 +1,45 @@
 [run]
 branch = True
+include =
+    ./cosipy/*
+    ./cosipy/*/*
+    ./cosipy/cpkernel/*
+
 omit =
+    setup.py
     */__init__.py
-    cosipy/tests/*
+    constants.py
+    COSIPY.py
+    config.py
+    slurm_config.py
     data/*
     docs/*
     logs/*
     utilities/hpc/*
+
+[html]
+# avoid workspace clutter
+directory = ./logs/coverage/
+
+[report]
+# Regexes for lines to exclude from consideration
+exclude_lines =
+    # Re-enable standard pragma
+    pragma: no cover
+
+    # Don't complain about missing debug-only code:
+    def __repr__
+    if self\.debug
+
+    # Don't complain if tests don't hit defensive assertion code:
+    raise AssertionError
+    raise NotImplementedError
+
+    # Don't complain if non-runnable code isn't run:
+    if 0:
+    if __name__ == .__main__.:
+
+    # Don't complain about abstract methods as they aren't run:
+    @(abc\.)?abstractmethod
+
+ignore_errors = True
diff --git a/.gitignore b/.gitignore
index f80601dd..26066f31 100644
--- a/.gitignore
+++ b/.gitignore
@@ -14,3 +14,4 @@ error*
 slurm.err
 slurm.out
 .coverage
+.vscode/
\ No newline at end of file
diff --git a/COSIPY.py b/COSIPY.py
index 9ec24fae..c9be3160 100644
--- a/COSIPY.py
+++ b/COSIPY.py
@@ -22,31 +22,30 @@
     Correspondence: tobias.sauter@fau.de
 
 """
+import cProfile
+import logging
 import os
+import sys
 from datetime import datetime
 from itertools import product
-import itertools
 
-import logging
+import numpy as np
+import pandas as pd
+import scipy
 import yaml
-
-from config import *
-from slurm_config import *
-from cosipy.cpkernel.cosipy_core import * 
-from cosipy.cpkernel.io import *
-
+# from dask import compute, delayed
+# import dask as da
+# from dask.diagnostics import ProgressBar
+from dask.distributed import as_completed, progress
+from dask_jobqueue import SLURMCluster
 from distributed import Client, LocalCluster
-from dask import compute, delayed
-import dask as da
-from dask.diagnostics import ProgressBar
-from dask.distributed import progress, wait, as_completed
-import dask
+# import dask
 from tornado import gen
-from dask_jobqueue import SLURMCluster
-
-import scipy
 
-import cProfile
+from config import *
+from cosipy.cpkernel.cosipy_core import cosipy_core
+from cosipy.cpkernel.io import IOClass
+from slurm_config import *
 
 def main():
 
@@ -217,13 +216,13 @@ def run_cosipy(cluster, IO, DATA, RESULT, RESTART, futures):
                 
             if WRF is True:
                 mask = DATA.MASK.sel(south_north=y, west_east=x)
-	        # Provide restart grid if necessary
-                if ((mask==1) & (restart==False)):
+                # Provide restart grid if necessary
+                if ((mask==1) & (not restart)):
                     if np.isnan(DATA.sel(south_north=y, west_east=x).to_array()).any():
                         print('ERROR!!!!!!!!!!! There are NaNs in the dataset')
                         sys.exit()
                     futures.append(client.submit(cosipy_core, DATA.sel(south_north=y, west_east=x), y, x, stake_names=stake_names, stake_data=df_stakes_data))
-                elif ((mask==1) & (restart==True)):
+                elif ((mask==1) & (restart)):
                     if np.isnan(DATA.sel(south_north=y, west_east=x).to_array()).any():
                         print('ERROR!!!!!!!!!!! There are NaNs in the dataset')
                         sys.exit()
@@ -232,13 +231,13 @@ def run_cosipy(cluster, IO, DATA, RESULT, RESTART, futures):
                                              stake_names=stake_names, stake_data=df_stakes_data))
             else:
                 mask = DATA.MASK.isel(lat=y, lon=x)
-	        # Provide restart grid if necessary
-                if ((mask==1) & (restart==False)):
+                # Provide restart grid if necessary
+                if ((mask==1) & (not restart)):
                     if np.isnan(DATA.isel(lat=y,lon=x).to_array()).any():
                         print('ERROR!!!!!!!!!!! There are NaNs in the dataset')
                         sys.exit()
                     futures.append(client.submit(cosipy_core, DATA.isel(lat=y, lon=x), y, x, stake_names=stake_names, stake_data=df_stakes_data))
-                elif ((mask==1) & (restart==True)):
+                elif ((mask==1) & (restart)):
                     if np.isnan(DATA.isel(lat=y,lon=x).to_array()).any():
                         print('ERROR!!!!!!!!!!! There are NaNs in the dataset')
                         sys.exit()
@@ -306,7 +305,7 @@ def run_cosipy(cluster, IO, DATA, RESULT, RESTART, futures):
 
 
 def start_logging():
-    ''' Start the python logging'''
+    """Start the python logging"""
 
     if os.path.exists('./cosipy.yaml'):
         with open('./cosipy.yaml', 'rt') as f:
@@ -327,7 +326,7 @@ def transform_coordinates(coords):
     A = 6378.137 # major axis [km]   
     E2 = 6.69437999014e-3 # eccentricity squared    
     
-    coords = np.asarray(coords).astype(np.float)
+    coords = np.asarray(coords).astype(float)
                                                       
     # is coords a tuple? Convert it to an one-element array of tuples
     if coords.ndim == 1:
@@ -361,6 +360,6 @@ def close_everything(scheduler):
 
 
 
-''' MODEL EXECUTION '''
+""" MODEL EXECUTION """
 if __name__ == "__main__":
     main()
diff --git a/constants.py b/constants.py
index 45340ec7..dca77ffd 100644
--- a/constants.py
+++ b/constants.py
@@ -16,7 +16,7 @@
 penetrating_method = 'Bintanja95'               # possibilities: 'Bintanja95'
 roughness_method = 'Moelg12'                    # possibilities: 'Moelg12'
 saturation_water_vapour_method = 'Sonntag90'    # possibilities: 'Sonntag90'
-thermal_conductivity_method = 'bulk'		    # possibilities: 'bulk', 'empirical'
+thermal_conductivity_method = 'bulk'		# possibilities: 'bulk', 'empirical'
 sfc_temperature_method = 'SLSQP'                # possibilities: 'L-BFGS-B', 'SLSQP'(faster), 'Newton' (Secant, fastest)'
 
 # WRF_X_CSPY: for efficiency and consistency
diff --git a/cosipy/cpkernel/cosipy_core.py b/cosipy/cpkernel/cosipy_core.py
index ea96e9d7..6f865e0d 100644
--- a/cosipy/cpkernel/cosipy_core.py
+++ b/cosipy/cpkernel/cosipy_core.py
@@ -210,9 +210,9 @@ def cosipy_core(DATA, indY, indX, GRID_RESTART=None, stake_names=None, stake_dat
 
         if SNOWFALL > 0.0:
             # Add a new snow node on top
-           GRID.add_fresh_snow(SNOWFALL, density_fresh_snow, np.minimum(float(T2[t]),zero_temperature), 0.0)
+            GRID.add_fresh_snow(SNOWFALL, density_fresh_snow, np.minimum(float(T2[t]),zero_temperature), 0.0)
         else:
-           GRID.set_fresh_snow_props_update_time(dt)
+            GRID.set_fresh_snow_props_update_time(dt)
 
         # Guarantee that solar radiation is greater equal zero
         if (G[t]<0.0):
diff --git a/cosipy/cpkernel/grid.py b/cosipy/cpkernel/grid.py
index 9ec5e375..b3486390 100644
--- a/cosipy/cpkernel/grid.py
+++ b/cosipy/cpkernel/grid.py
@@ -1,53 +1,61 @@
-import numpy as np
-from constants import *
-from config import *
-from cosipy.cpkernel.node import *
 import os
-
 from collections import OrderedDict
-from numba import types, typed, intp, float64, optional
+
+import numpy as np
+from numba import float64, intp, optional, typed, types
 from numba.experimental import jitclass
 
+import constants
+from cosipy.cpkernel.node import Node
+
 node_type = Node.class_type.instance_type
 
 spec = OrderedDict()
-spec['layer_heights'] = float64[:]
-spec['layer_densities'] = float64[:]
-spec['layer_temperatures'] = float64[:]
-spec['layer_liquid_water_content'] = float64[:]
-spec['layer_ice_fraction'] = optional(float64[:])
-spec['number_nodes'] = intp
-spec['new_snow_height'] = float64
-spec['new_snow_timestamp'] = float64
-spec['old_snow_timestamp'] = float64
-spec['grid'] = types.ListType(node_type)     
+spec["layer_heights"] = float64[:]
+spec["layer_densities"] = float64[:]
+spec["layer_temperatures"] = float64[:]
+spec["layer_liquid_water_content"] = float64[:]
+spec["layer_ice_fraction"] = optional(float64[:])
+spec["number_nodes"] = intp
+spec["new_snow_height"] = float64
+spec["new_snow_timestamp"] = float64
+spec["old_snow_timestamp"] = float64
+spec["grid"] = types.ListType(node_type)
 
 @jitclass(spec)
 class Grid:
-
-    def __init__(self, layer_heights, layer_densities, layer_temperatures, layer_liquid_water_content, layer_ice_fraction=None,
-        new_snow_height=None, new_snow_timestamp=None, old_snow_timestamp=None):
-        """ The Grid-class controls the numerical mesh. 
-        
-        The grid consists of a list of nodes (layers) that store the information 
-        of individual layers. The class provides various setter/getter functions
-        to add, read, overwrite, merge, split, update or re-mesh the layers. 
-
-        Parameters
+    def __init__(
+        self,
+        layer_heights,
+        layer_densities,
+        layer_temperatures,
+        layer_liquid_water_content,
+        layer_ice_fraction=None,
+        new_snow_height=None,
+        new_snow_timestamp=None,
+        old_snow_timestamp=None,
+    ):
+        """The Grid-class controls the numerical mesh.
+
+        The grid consists of a list of nodes (layers) that store the
+        information of individual layers. The class provides various
+        setter/getter functions to add, read, overwrite, merge,
+        split, update or re-mesh the layers.
+
+        Attributes
         ----------
-            layer_heights : float
-                Height of the snowpack layers [:math:`m`]
-            layer_densities : float
-                Snow density of the snowpack layers [:math:`kg~m^{-3}`]
-            layer_temperatures : float
-                Temperatures of the layers [:math:`K`]
-            layer_liquid_water_content : float
-                Liquid water content of the layers [:math:`m~w.e.`]
+        layer_heights : np.ndarray
+            Height of the snowpack layers [:math:`m`].
+        layer_densities : np.ndarray
+            Snow density of the snowpack layers [:math:`kg~m^{-3}`].
+        layer_temperatures : np.ndarray
+            Layer temperatures [:math:`K`].
+        layer_liquid_water_content : np.ndarray
+            Liquid water content of the layers [:math:`m~w.e.`].
 
         Returns
         -------
-            Node : :py:class:`cosipy.cpkernel.grid` object
-
+        Grid : :py:class:`cosipy.cpkernel.grid` object.
         """
         # Set class variables
         self.layer_heights = layer_heights
@@ -61,15 +69,22 @@ def __init__(self, layer_heights, layer_densities, layer_temperatures, layer_liq
 
         # Track the fresh snow layer (new_snow_height, new_snow_timestamp) as well as the old
         # snow layer age (old_snow_timestamp)
-        if (new_snow_height is not None) and (new_snow_timestamp is not None) and \
-           (old_snow_timestamp is not None):
-            self.new_snow_height = new_snow_height         # meter snow accumulation
-            self.new_snow_timestamp = new_snow_timestamp   # seconds since snowfall
-            self.old_snow_timestamp = old_snow_timestamp   # snow age below fresh snow layer
+        if (
+            (new_snow_height is not None)
+            and (new_snow_timestamp is not None)
+            and (old_snow_timestamp is not None)
+        ):
+            self.new_snow_height = new_snow_height  # meter snow accumulation
+            self.new_snow_timestamp = (
+                new_snow_timestamp  # seconds since snowfall
+            )
+            self.old_snow_timestamp = (
+                old_snow_timestamp  # snow age below fresh snow layer
+            )
         else:
-	    #TO DO: pick better initialization values
-            self.new_snow_height = 0.0      
-            self.new_snow_timestamp = 0.0   
+            # TO DO: pick better initialization values
+            self.new_snow_height = 0.0
+            self.new_snow_timestamp = 0.0
             self.old_snow_timestamp = 0.0
 
         # Do the grid initialization
@@ -77,140 +92,183 @@ def __init__(self, layer_heights, layer_densities, layer_temperatures, layer_liq
 
         self.init_grid()
 
-    
     def init_grid(self):
-        """ Initialize the grid with according to the input data """
+        """Initialize the grid with according to the input data"""
         # Fill the list with node instances and fill it with user defined data
         for idxNode in range(self.number_nodes):
             layer_IF = None
             if self.layer_ice_fraction is not None:
                 layer_IF = self.layer_ice_fraction[idxNode]
-            self.grid.append(Node(self.layer_heights[idxNode], self.layer_densities[idxNode],
-                self.layer_temperatures[idxNode], self.layer_liquid_water_content[idxNode], layer_IF))
-
-
-    def add_fresh_snow(self, height, density, temperature, liquid_water_content):
-        """ Adds a fresh snow layer (node) at the beginning of the node list (upper layer) 
+            self.grid.append(
+                Node(
+                    self.layer_heights[idxNode],
+                    self.layer_densities[idxNode],
+                    self.layer_temperatures[idxNode],
+                    self.layer_liquid_water_content[idxNode],
+                    layer_IF,
+                )
+            )
+
+    def add_fresh_snow(
+        self, height, density, temperature, liquid_water_content
+    ):
+        """Add a fresh snow layer (node).
+         
+        Adds a fresh snow layer to the beginning of the node list (upper
+        layer).
 
         Parameters
         ----------
-            height : float
-                Height of the layer [:math:`m`]
-            density : float
-                Density of the layer [:math:`kg~m^{-3}`]
-            temperature : float
-                Temperature of the layer [:math:`K`]
-            liquid_water_content : float
-                Liquid water content of the layer [:math:`m~w.e.`]
+        height : float
+            Layer height [:math:`m`].
+        density : float
+            Layer density [:math:`kg~m^{-3}`].
+        temperature : float
+            Layer temperature [:math:`K`].
+        liquid_water_content : float
+            Liquid water content of the layer [:math:`m~w.e.`].
         """
-	
+
         # Add new node
-        self.grid.insert(0, Node(height, density, temperature, liquid_water_content, None))
+        self.grid.insert(
+            0, Node(height, density, temperature, liquid_water_content, None)
+        )
 
         # Increase node counter
         self.number_nodes += 1
 
-        # Set the fresh snow properties for albedo calculation (height and timestamp)
+        # Set fresh snow properties for albedo calculation (height and
+        # timestamp)
         self.set_fresh_snow_props(height)
 
+    def remove_node(self, idx: list = None):
+        """Remove a layer (node) from the grid (node list).
 
-
-    def remove_node(self, idx=None):
-        """ Removes a layer (node) from the grid (node list). 
-        
-        Parameters
-        ----------
-            idx : int
-                Index of the node to be removed. If no index is provided the first
-                node is beeing removed.
-
+        Args:
+            idx: Indices of the node to be removed. If empty or None,
+                the first node is removed.
         """
 
         # Remove node from list when there is at least one node
-        if not self.grid:
-            pass
-        else:
-            if idx is None:
+        if self.grid:
+            if idx is None or not idx:
                 self.grid.pop(0)
+                self.number_nodes -= 1  # Decrease node counter
             else:
                 for index in sorted(idx, reverse=True):
                     del self.grid[index]
-
-            # Decrease node counter
-            self.number_nodes -= len(idx)
-
+                self.number_nodes -= len(idx)
 
     def merge_nodes(self, idx):
-        """ Merges two subsequent nodes. 
-        
-        This function merges the two nodes at location idx and idx+1. The node at idx is updated
-        with the new properties (height, liquid water content, ice fraction, temperature), while the node
-        at idx+1 is deleted after merging.
+        """Merge two subsequent nodes.
+
+        Merges the two nodes at location `idx` and `idx+1`.
+        The node at `idx` is updated with the new properties (height,
+        liquid water content, ice fraction, temperature). The node
+        at `idx+1` is deleted after merging.
 
         Parameters
         ----------
-            idx : int
-                Index of the node to be removed. If no index is provided the first
-                node is beeing removed.
-        
-        """ 
+        idx : int
+            Index of the node to be removed. The first node is removed
+            if no index is provided.
+        """
         # Get overburden pressure for consistency check
-        w0 = self.get_node_height(idx)*self.get_node_density(idx)+self.get_node_height(idx+1)*self.get_node_density(idx+1)
+        w0 = self.get_node_height(idx) * self.get_node_density(
+            idx
+        ) + self.get_node_height(idx + 1) * self.get_node_density(idx + 1)
 
         # New layer height by adding up the height of the two layers
-        new_height = self.get_node_height(idx) + self.get_node_height(idx+1)
+        new_height = self.get_node_height(idx) + self.get_node_height(idx + 1)
 
         # Update liquid water
-        #new_liquid_water_content = self.get_node_liquid_water_content(idx) + self.get_node_liquid_water_content(idx+1)
-        new_liquid_water_content = (self.get_node_liquid_water_content(idx) * self.get_node_height(idx) + \
-                                    self.get_node_liquid_water_content(idx+1) * self.get_node_height(idx + 1))/new_height
+        # new_liquid_water_content = self.get_node_liquid_water_content(idx) + self.get_node_liquid_water_content(idx+1)
+        new_liquid_water_content = (
+            self.get_node_liquid_water_content(idx) * self.get_node_height(idx)
+            + self.get_node_liquid_water_content(idx + 1)
+            * self.get_node_height(idx + 1)
+        ) / new_height
 
         # Update ice fraction
-        new_ice_fraction = ((self.get_node_ice_fraction(idx)*self.get_node_height(idx) + \
-                            self.get_node_ice_fraction(idx+1)*self.get_node_height(idx+1))/new_height)
+        new_ice_fraction = (
+            self.get_node_ice_fraction(idx) * self.get_node_height(idx)
+            + self.get_node_ice_fraction(idx + 1)
+            * self.get_node_height(idx + 1)
+        ) / new_height
 
         # New air porosity
         new_air_porosity = 1 - new_liquid_water_content - new_ice_fraction
 
-        if abs(1-new_ice_fraction-new_air_porosity-new_liquid_water_content)>1e-8:
-            print('Merging is not mass consistent',(new_ice_fraction+new_air_porosity+new_liquid_water_content))
+        if (
+            abs(
+                1
+                - new_ice_fraction
+                - new_air_porosity
+                - new_liquid_water_content
+            )
+            > 1e-8
+        ):
+            print(
+                "Merging is not mass consistent",
+                (
+                    new_ice_fraction
+                    + new_air_porosity
+                    + new_liquid_water_content
+                ),
+            )
 
         # Calc new temperature
-        new_temperature = (self.get_node_height(idx)/new_height)*self.get_node_temperature(idx) + \
-                            (self.get_node_height(idx+1)/new_height)*self.get_node_temperature(idx+1)
+        new_temperature = (
+            self.get_node_height(idx) / new_height
+        ) * self.get_node_temperature(idx) + (
+            self.get_node_height(idx + 1) / new_height
+        ) * self.get_node_temperature(
+            idx + 1
+        )
 
         # Update node properties
-        self.update_node(idx, new_height, new_temperature, new_ice_fraction, new_liquid_water_content)
-        
-        # Remove the second layer
-        self.remove_node([idx+1])
-
+        self.update_node(
+            idx,
+            new_height,
+            new_temperature,
+            new_ice_fraction,
+            new_liquid_water_content,
+        )
 
+        # Remove the second layer
+        self.remove_node([idx + 1])
 
     def correct_layer(self, idx, min_height):
-        """ Adjusts the height of a given layer.
+        """Adjust the height of a given layer.
+
+        Adjusts the height of the layer at index `idx` to the given
+        height `min_height`. First the layers below are merged until the
+        height is sufficiently large to allow for the adjustment. Then
+        the layer is merged with the subsequent layer.
 
-        This function adjusts the height of the layer at index idx to the given height min_height. 
-        First, the layers below are merged until the height is sufficiently large to allow the 
-        adjustment. Then the layer is merged with the subsequent layer. 
-        
         Parameters
         ----------
-            idx : int
-                Index of the node to be removed. If no index is provided the first
-                node is beeing removed.
-            min_height : float
-                New layer height [:math:`m`]
-        
+        idx : int
+            Index of the node to be removed. The first node is removed
+            if no index is provided.
+        min_height : float
+            New layer height [:math:`m`].
+
         """
         # New layer height by adding up the height of the two layers
         total_height = self.get_node_height(idx)
 
         # Merge subsequent layer with underlying layers until height of the layer is greater than the given height
-        while ((total_height<min_height) & (idx+1<self.get_number_layers())):
-            if (self.get_node_density(idx)<snow_ice_threshold) & (self.get_node_density(idx+1)<snow_ice_threshold):
+        while (total_height < min_height) & (
+            idx + 1 < self.get_number_layers()
+        ):
+            if (self.get_node_density(idx) < constants.snow_ice_threshold) & (
+                self.get_node_density(idx + 1) < constants.snow_ice_threshold
+            ):
                 self.merge_nodes(idx)
-            elif (self.get_node_density(idx)>=snow_ice_threshold) & (self.get_node_density(idx+1)>=snow_ice_threshold):
+            elif (self.get_node_density(idx) >= constants.snow_ice_threshold) & (
+                self.get_node_density(idx + 1) >= constants.snow_ice_threshold
+            ):
                 self.merge_nodes(idx)
             else:
                 break
@@ -219,23 +277,39 @@ def correct_layer(self, idx, min_height):
             total_height = self.get_node_height(idx)
 
         # Only merge snow-snow or glacier-glacier, and if the height is greater than the minimum height
-        if (total_height>min_height):
-
+        if total_height > min_height:
             # Get new heights for layer 0 and 1
             h0 = min_height
             h1 = total_height - min_height
 
-            # Update liquid water
-            # First the upper layer is filled with water until the max.
-            # retention. The rest is assigned to the second layer.
-            # if LWC exceeds the irreducible water content of the first layer,
-            # then fill the first layer and assign the rest to the second layer
-            if ((self.get_node_liquid_water_content(idx)-self.get_node_irreducible_water_content(idx))>0):
-                lw0 = self.get_node_irreducible_water_content(idx) * self.get_node_height(idx)
-                lw1 = self.get_node_liquid_water_content(idx) * self.get_node_height(idx) - self.get_node_irreducible_water_content(idx) * self.get_node_height(idx)
+            """Update liquid water content.
+            Fills the upper layer with water until maximum retention.
+            The remaining water is assigned to the second layer.
+            If LWC exceeds the irreducible water content of the first
+            layer, then the first layer is filled and the rest assigned
+            to the second layer.
+            """
+            if (
+                self.get_node_liquid_water_content(idx)
+                - self.get_node_irreducible_water_content(idx)
+            ) > 0:
+                lw0 = self.get_node_irreducible_water_content(
+                    idx
+                ) * self.get_node_height(idx)
+                lw1 = self.get_node_liquid_water_content(
+                    idx
+                ) * self.get_node_height(
+                    idx
+                ) - self.get_node_irreducible_water_content(
+                    idx
+                ) * self.get_node_height(
+                    idx
+                )
             # if LWC<WC_irr, then assign all water to the first layer
-            else:   
-                lw0 = self.get_node_liquid_water_content(idx) * self.get_node_height(idx)
+            else:
+                lw0 = self.get_node_liquid_water_content(
+                    idx
+                ) * self.get_node_height(idx)
                 lw1 = 0.0
 
             # Update ice fraction
@@ -247,34 +321,47 @@ def correct_layer(self, idx, min_height):
             T1 = self.get_node_temperature(idx)
 
             # New volume fractions and density
-            lwc0 = lw0/h0
-            lwc1 = lw1/h1
+            lwc0 = lw0 / h0
+            lwc1 = lw1 / h1
             por0 = 1 - lwc0 - if0
             por1 = 1 - lwc1 - if1
 
             # Check for consistency
-            if (abs(1-if0-por0-lwc0)>1e-8) | (abs(1-if1-por1-lwc1)>1e-8):
-                print('Correct layer is not mass consistent [Layer 0]',(if0,por0,lwc0))
-                print('Correct layer is not mass consistent [Layer 1]',(if0,por0,lwc0))
+            if (abs(1 - if0 - por0 - lwc0) > 1e-8) | (
+                abs(1 - if1 - por1 - lwc1) > 1e-8
+            ):
+                print(
+                    "Correct layer is not mass consistent [Layer 0]",
+                    (if0, por0, lwc0),
+                )
+                print(
+                    "Correct layer is not mass consistent [Layer 1]",
+                    (if0, por0, lwc0),
+                )
 
             # Update node properties
             self.update_node(idx, h0, T0, if0, lwc0)
-            self.grid.insert(idx+1, Node(h1, self.get_node_density(idx), T1, lwc1, if1))
+            self.grid.insert(
+                idx + 1, Node(h1, self.get_node_density(idx), T1, lwc1, if1)
+            )
 
             # Update node counter
             self.number_nodes += 1
 
-
     def log_profile(self):
-        """ Logarithmic re-meshing. 
-        
-        The logirithmic algorithm re-meshes the layer heights (numerical grid) using a stretching
-        factor and a given first layer height. The latter one is provided by the first_layer_height
-        constant defined in the config.py. The layer_stretching variable defines the streching
-        factor, e.g. 1.1 corresponds to a 10% streching from one layer to the next.
+        """Remesh the layer heights logarithmically.
+
+        This algorithm remeshes the layer heights (numerical
+        grid) logarithmically using a given stretching factor and first
+        layer height. Both are defined in `constants.py`:
+
+        * The stretching factor is defined by `layer_stretching`.
+        * The first layer height is defined by `first_layer_height`.
         
+        E.g. for the stretching factor, a value of 1.1 corresponds to a
+        10% stretching from one layer to the next.
         """
-        last_layer_height = first_layer_height
+        last_layer_height = constants.first_layer_height
 
         # Total snowheight
         hsnow = self.get_total_snowheight()
@@ -282,262 +369,296 @@ def log_profile(self):
         # How much snow is not remeshed
         hrest = hsnow
 
-        # First, the snowpack is remeshed
+        # First remesh the snowpack
         idx = 0
 
-
-        while (idx < self.get_number_snow_layers()):
-
-            if (hrest>=last_layer_height):
+        while idx < self.get_number_snow_layers():
+            if hrest >= last_layer_height:
                 # Correct first layer
-                self.correct_layer(idx,last_layer_height)
+                self.correct_layer(idx, last_layer_height)
 
                 hrest = hrest - last_layer_height
 
                 # Height for the next layer
-                last_layer_height = layer_stretching*last_layer_height
-
-            # if the last layer is smaller than the required height, then merge
-            # with the previous layer
-            elif ((hrest<last_layer_height) & (idx>0)):
-                self.merge_nodes(idx-1)
+                last_layer_height = constants.layer_stretching * last_layer_height
 
-            idx = idx+1
+            # if the last layer is smaller than the required height,
+            # then merge with the previous layer
+            elif (hrest < last_layer_height) & (idx > 0):
+                self.merge_nodes(idx - 1)
 
+            idx = idx + 1
 
-        # get the glacier depth 
-        hrest = self.get_total_height()-self.get_total_snowheight()
+        # get the glacier depth
+        hrest = self.get_total_height() - self.get_total_snowheight()
 
         # get number of snow layers
         idx = self.get_number_snow_layers()
 
         # then the glacier
-        while (idx < self.get_number_layers()):
-
-            if (hrest>=last_layer_height):
+        while idx < self.get_number_layers():
+            if hrest >= last_layer_height:
                 # Correct first layer
-                self.correct_layer(idx,last_layer_height)
+                self.correct_layer(idx, last_layer_height)
 
                 hrest = hrest - last_layer_height
 
                 # Height for the next layer
-                last_layer_height = layer_stretching*last_layer_height
+                last_layer_height = constants.layer_stretching * last_layer_height
 
-            # if the last layer is smaller than the required height, then merge
-            # with the previous layer
-            elif ((hrest<last_layer_height)):
-                self.merge_nodes(idx-1)
+            # if the last layer is smaller than the required height,
+            # then merge with the previous layer
+            elif hrest < last_layer_height:
+                self.merge_nodes(idx - 1)
 
-            idx = idx+1
+            idx = idx + 1
 
+    def adaptive_profile(self):
+        """Remesh according to certain layer state criteria.
 
+        This algorithm is an alternative to logarithmic remeshing.
+        It checks the similarity of two subsequent layers. Layers are
+        merged if:
 
-    def adaptive_profile(self):
-        """ Remesh according to certain layer state criteria.
-        
-        This algorithm is an alternative to the logarithmic re-meshing. It checks for similarity of
-        two subsequent layers. Layers are merged, if:
+        (1) the density difference between the layer and the subsequent
+        layer is smaller than the user defined threshold.
+        (2) the temperature difference is smaller than the user defined
+        threshold.
 
-        (1) the density difference between the layer and the subsequent layer is smaller than the user defined threshold
-        (2) the temperature difference is smaller than the user defined threshold
-        
-        The temperature_threshold_merging and density_threshold_merging variables in the
-        configuration file (config.py) define the corresponding thresholds. 
+        The thresholds are defined by `temperature_threshold_merging`
+        and `density_threshold_merging` in `constants.py`.
         """
-        # First, the snowpack is remeshed
+        # First remesh the snowpack
         idx = 0
         merge_counter = 0
-        while ((idx < self.get_number_snow_layers()-1)):
-
-            dT = np.abs(self.get_node_temperature(idx)-self.get_node_temperature(idx+1))
-            dRho = np.abs(self.get_node_density(idx)-self.get_node_density(idx+1))
-
-            if ((dT<=temperature_threshold_merging) & (dRho<=density_threshold_merging) & (self.get_node_height(idx)<=0.1) & (merge_counter<=merge_max)):
+        while idx < self.get_number_snow_layers() - 1:
+            dT = np.abs(
+                self.get_node_temperature(idx)
+                - self.get_node_temperature(idx + 1)
+            )
+            dRho = np.abs(
+                self.get_node_density(idx) - self.get_node_density(idx + 1)
+            )
+
+            if (
+                (dT <= constants.temperature_threshold_merging)
+                & (dRho <= constants.density_threshold_merging)
+                & (self.get_node_height(idx) <= 0.1)
+                & (merge_counter <= constants.merge_max)
+            ):
                 self.merge_nodes(idx)
                 merge_counter = merge_counter + 1
-            #elif ((self.get_node_height(idx)<=minimum_snow_layer_height) & (dRho<=density_threshold_merging)):
-            elif ((self.get_node_height(idx)<=minimum_snow_layer_height)):
+            # elif ((self.get_node_height(idx)<=minimum_snow_layer_height) & (dRho<=density_threshold_merging)):
+            elif self.get_node_height(idx) <= constants.minimum_snow_layer_height:
                 self.remove_node([idx])
             else:
                 idx += 1
 
         # Correct first layer
-        self.correct_layer(0 ,first_layer_height)
-
-
+        self.correct_layer(0, constants.first_layer_height)
 
     def split_node(self, pos):
-        """ Split node at position pos 
+        """Split node at position.
 
-        This function splits a node at location idx into two similar nodes. The new nodes at
-        location idx and idx+1 will have the same properties  (height, liquid water content, 
-        ice fraction, temperature).
+        Splits a node at a location index `pos` into two similar nodes.
+        The new nodes at location `pos` and `pos+1` will have the same
+        properties (height, liquid water content, ice fraction,
+        temperature).
 
         Parameters
         ----------
-            idx : int
-                Index of the node to be splitted.        
+        pos : int
+            Index of the node to split.
         """
-        self.grid.insert(pos+1, Node(self.get_node_height(pos)/2.0, self.get_node_density(pos), self.get_node_temperature(pos), \
-                                     self.get_node_liquid_water_content(pos)/2.0, self.get_node_ice_fraction(pos)))
-        self.update_node(pos, self.get_node_height(pos)/2.0, self.get_node_temperature(pos), \
-                                     self.get_node_ice_fraction(pos), self.get_node_liquid_water_content(pos)/2.0)
+        self.grid.insert(
+            pos + 1,
+            Node(
+                self.get_node_height(pos) / 2.0,
+                self.get_node_density(pos),
+                self.get_node_temperature(pos),
+                self.get_node_liquid_water_content(pos) / 2.0,
+                self.get_node_ice_fraction(pos),
+            ),
+        )
+        self.update_node(
+            pos,
+            self.get_node_height(pos) / 2.0,
+            self.get_node_temperature(pos),
+            self.get_node_ice_fraction(pos),
+            self.get_node_liquid_water_content(pos) / 2.0,
+        )
 
         self.number_nodes += 1
 
+    def update_node(
+        self, idx, height, temperature, ice_fraction, liquid_water_content
+    ):
+        """Update properties of a specific node.
 
-
-    def update_node(self, idx, height, temperature, ice_fraction, liquid_water_content):
-        """ Update properties of a specific node.
-
-        This function updates the properties height, temperature, ice_fraction and liquid water
-        content of a layer. The density cannot be updated since it is derived from air porosity,
-        liquid water content and ice fraction.
+        This function updates a layer's attributes for `height`,
+        `temperature`, `ice_fraction`, and `liquid_water_content`.
+        The density cannot be updated as it is derived from air
+        porosity, liquid water content, and ice fraction.
 
         Parameters
         ----------
-            idx : int
-                Index of the layer to be updated.
-            layer_heights : float
-                New height of the snowpack layers [:math:`m`]
-            layer_temperatures : float
-                New temperatures of the layers [:math:`K`]
-            ice_fraction : float
-                New ice fraction of the layer [:math:`-`]
-            layer_liquid_water_content : float
-                New liquid water content of the layers [:math:`m~w.e.`]
-        
+        idx : int
+            Index of the layer to be updated.
+        height : float
+            Layer's new snowpack height [:math:`m`].
+        temperature : float
+            Layer's new temperature [:math:`K`].
+        ice_fraction : float
+            Layer's new ice fraction [:math:`-`].
+        liquid_water_content : float
+            Layer's new liquid water content [:math:`m~w.e.`].
         """
-        self.set_node_height(idx,height)
-        self.set_node_temperature(idx,temperature)
-        self.set_node_ice_fraction(idx,ice_fraction)
-        self.set_node_liquid_water_content(idx,liquid_water_content)
-
-
+        self.set_node_height(idx, height)
+        self.set_node_temperature(idx, temperature)
+        self.set_node_ice_fraction(idx, ice_fraction)
+        self.set_node_liquid_water_content(idx, liquid_water_content)
 
     def check(self, name):
-        """ Function checks whether temperature and layer heights are within the valid range """
+        """Check layer temperature and height are within a valid range."""
         if np.min(self.get_height()) < 0.01:
             print(name)
-            print('Layer height is smaller than the user defined minimum new_height')
+            print(
+                "Layer height is smaller than the user defined minimum new_height"
+            )
             print(self.get_height())
             print(self.get_density())
         if np.max(self.get_temperature()) > 273.2:
             print(name)
-            print('Layer temperature exceeds 273.16 K')
+            print("Layer temperature exceeds 273.16 K")
             print(self.get_temperature())
             print(self.get_density())
         if np.max(self.get_height()) > 1.0:
             print(name)
-            print('Layer height exceeds 1.0 m')
+            print("Layer height exceeds 1.0 m")
             print(self.get_height())
             print(self.get_density())
 
-
-
     def update_grid(self):
-        """ Re-meshes the layers (numerical grid).
+        """Remesh the layers (numerical grid).
 
-            Two algorithms are currently implemented to re-mesh the layers:
+        Two algorithms are currently implemented to remesh layers:
 
-                (i)  log_profile
-                (ii) adaptive_profile
+            (i)  log_profile
+            (ii) adaptive_profile
 
-            (i)  The log-profile algorithm arranges the mesh logarithmically.
-                 The user provides a stretching factor (layer_stretching in the configuration file) 
-                 that determines the increase in layer heights.
-
-            (ii) The adjustment of the profile is done on the basis of the similarity of layers. 
-                Layers with very similar states (temperature and density) are joined together. The
-                 similarity is determined by user-specified threshold values
-                 (temperature_threshold_merging, density_threshold_merging). In
-                 addition, the maximum number of merging steps per time step
-                 can be specified (merge_max).
+        (i)  The log-profile algorithm arranges the mesh
+             logarithmically. The user specifies the stretching factor
+             `layer_stretching` in `constants.py` to determine the
+             increase in layer heights.
 
+        (ii) Profile adjustment is done using layer similarity. Layers
+             with very similar temperature and density states are
+             joined. Similarity is determined from the user-specified
+             threshold values `temperature_threshold_merging` and
+             `density_threshold_merging` in `constants.py`. The maximum
+             number of merging steps per time step is specified by
+             `merge_max`.
         """
-        #-------------------------------------------------------------------------
+        # -------------------------------------------------------------------------
         # Remeshing options
-        #-------------------------------------------------------------------------
-        if (remesh_method=='log_profile'):
+        # -------------------------------------------------------------------------
+        if constants.remesh_method == "log_profile":
             self.log_profile()
-        elif (remesh_method=='adaptive_profile'):
+        elif constants.remesh_method == "adaptive_profile":
             self.adaptive_profile()
 
         # if first layer becomes very small, remove it
-        if (self.get_node_height(0)<minimum_snow_layer_height):
+        if self.get_node_height(0) < constants.minimum_snow_layer_height:
             self.remove_node([0])
 
-
     def merge_snow_with_glacier(self, idx):
-        """ Merges a snow layer with a ice layer.
-        
-        The function merges a snow layer at location idx (density smaller than the snow_ice_threshold value) with an
-        ice layer at location idx+1.
+        """Merge a snow layer with an ice layer.
+
+        Merges a snow layer at location `idx` (density smaller than the
+        `snow_ice_threshold` value in `constants.py`) with an ice layer
+        at location `idx+1`.
 
         Parameters
         ----------
-            idx : int
-                Index of the snow layer.
+        idx : int
+            Index of the snow layer.
         """
-        if (self.get_node_density(idx) < snow_ice_threshold) & (self.get_node_density(idx+1) >= snow_ice_threshold):
-
+        if (self.get_node_density(idx) < constants.snow_ice_threshold) & (
+            self.get_node_density(idx + 1) >= constants.snow_ice_threshold
+        ):
             # Update node properties
-            first_layer_height = self.get_node_height(idx)*(self.get_node_density(idx)/ice_density)
-            self.update_node(idx+1, self.get_node_height(idx+1)+first_layer_height, self.get_node_temperature(idx+1), self.get_node_ice_fraction(idx+1), 0.0)
+            first_layer_height = self.get_node_height(idx) * (
+                self.get_node_density(idx) / constants.ice_density
+            )
+            self.update_node(
+                idx + 1,
+                self.get_node_height(idx + 1) + first_layer_height,
+                self.get_node_temperature(idx + 1),
+                self.get_node_ice_fraction(idx + 1),
+                0.0,
+            )
 
             # Remove the second layer
             self.remove_node([idx])
 
-            #self.check('Merge snow with glacier function')
+            # self.check('Merge snow with glacier function')
 
+    def remove_melt_weq(self, melt, idx=0):
+        """Remove mass from a layer.
 
+        Reduces the mass/height of layer `idx` by the available melt
+        energy.
 
-    def remove_melt_weq(self, melt, idx=0):
-        """ Removes mass from a layer.
-        
-        The mass/height of layer idx is reduced by the available melt energy.
-        
         Parameters
         ----------
-            melt : float
-                Snow water equivalent of melt [:math:`m~w.e.`]
-            idx : int
-                Index of the layer. If no values is given, the function acts on the first
-                layer.
+        melt : float
+            Snow water equivalent of melt [:math:`m~w.e.`].
+        idx : int
+            Index of the layer. If no value is given, the function acts
+            on the first layer.
         """
         lwc_from_layers = 0
 
-        while melt>0:
+        while melt > 0:
             # Get SWE of layer
-            SWE = self.get_node_height(idx) * (self.get_node_density(idx)/water_density)
+            SWE = self.get_node_height(idx) * (
+                self.get_node_density(idx) / constants.water_density
+            )
             # Remove melt from layer and set new snowheight
-            if (melt<SWE):
-                self.set_node_height(idx, (SWE-melt)/(self.get_node_density(idx)/water_density))
+            if melt < SWE:
+                self.set_node_height(
+                    idx,
+                    (SWE - melt)
+                    / (self.get_node_density(idx) / constants.water_density),
+                )
                 melt = 0.0
             # remove layer otherwise and continue loop
-            elif (melt>=SWE):
-                lwc_from_layers = lwc_from_layers + self.get_node_liquid_water_content(idx) * self.get_node_height(idx)
+            elif melt >= SWE:
+                lwc_from_layers = (
+                    lwc_from_layers
+                    + self.get_node_liquid_water_content(idx)
+                    * self.get_node_height(idx)
+                )
                 self.remove_node([idx])
-                melt = melt-SWE
+                melt = melt - SWE
 
         # Keep track of the fresh snow layer
-        if (idx==0):
-            self.set_fresh_snow_props_height(self.new_snow_height-melt)
+        if idx == 0:
+            self.set_fresh_snow_props_height(self.new_snow_height - melt)
 
         return lwc_from_layers
 
-
-    #===============================================================================
+    # ===============================================================================
     # Getter and setter functions
-    #===============================================================================
+    # ===============================================================================
 
     def set_fresh_snow_props(self, height):
-        """ Keeps track of the new snowheight.
-        
+        """Track the new snowheight.
+
         Parameters
         ----------
-            height : float
-                Height of the fresh snow layer [:math:`m`].
+        height : float
+            Height of the fresh snow layer [:math:`m`].
         """
         self.new_snow_height = height
         # Keep track of the old snow age
@@ -546,351 +667,384 @@ def set_fresh_snow_props(self, height):
         self.new_snow_timestamp = 0
 
     def set_fresh_snow_props_to_old_props(self):
-        """ Sets the timestamp of the fresh snow properties back to the timestamp of the underlying snow layer.
-        
-        The function is used internally to keep track of the albedo properties of the first snow
-        layer.
+        """Revert the timestamp of fresh snow properties.
+
+        Reverts the timestamp of fresh snow properties to that of the
+        underlying snow layer. This is used internally to track the
+        albedo properties of the first snow layer.
         """
         self.new_snow_timestamp = self.old_snow_timestamp
 
     def set_fresh_snow_props_update_time(self, seconds):
-        """ Update timestamp of snow props.
+        """Update the timestamp of the snow properties.
 
         Parameters
         ----------
-            height : float
-                Height of the fresh snow layer [:math:`m`].
-            seconds : float
-                seconds without snowfall
-                [:math:`s`].
+        seconds : float
+            seconds without snowfall [:math:`s`].
         """
         self.old_snow_timestamp = self.old_snow_timestamp + seconds
         # Set the timestamp to zero
         self.new_snow_timestamp = self.new_snow_timestamp + seconds
 
     def set_fresh_snow_props_height(self, height):
-        """ Updates the fresh snow layer height property.
-        
-        The function is used internally to keep track of the albedo properties oir the first snow
-        layer.
+        """Update the fresh snow layer height.
+
+        This is used internally to track the albedo properties of the
+        first snow layer.
         """
         self.new_snow_height = height
-	
+
     def get_fresh_snow_props(self):
-        """ Returns the properties of the first snow layer.
+        """Get the first snow layer's properties.
 
-        The function is used internally to keep track of the albedo properties oir the first snow
-        layer.
+        This is used internally to track the albedo properties of the
+        first snow layer.
         """
-        return self.new_snow_height, self.new_snow_timestamp, self.old_snow_timestamp
-
+        return (
+            self.new_snow_height,
+            self.new_snow_timestamp,
+            self.old_snow_timestamp,
+        )
 
     def set_node_temperature(self, idx, temperature):
-        """ Sets the temperature of layer (node) at location idx.
-        
+        """Set the temperature of a layer (node) at location `idx`.
+
         Parameters
         ----------
-            idx : int
-                Index of the layer.
-            temperature : float
-                New layer temperature [:math:`K`].
+        idx : int
+            Index of the layer.
+        temperature : float
+            Layer's new temperature [:math:`K`].
         """
         self.grid[idx].set_layer_temperature(temperature)
 
-
-
     def set_temperature(self, temperature):
-        """ Sets the temperature of layer (node) at location idx.
-        
+        """Set all layer temperatures.
+
         Parameters
         ----------
-            idx : int
-                Index of the layer.
-            temperature : float
-                New layer temperature [:math:`K`].
+        temperature : np.ndarray
+            New layer temperatures [:math:`K`].
         """
         for idx in range(self.number_nodes):
             self.grid[idx].set_layer_temperature(temperature[idx])
 
-
-
-    def set_node_height(self, idx, height):
-        """ Set height of node idx """
+    def set_node_height(self, idx: int, height: float):
+        """Set node height."""
         self.grid[idx].set_layer_height(height)
 
-
-
-    def set_height(self, height):
-        """ Set height of profile """
+    def set_height(self, height: np.ndarray):
+        """Set the height profile."""
         for idx in range(self.number_nodes):
             self.grid[idx].set_layer_height(height[idx])
 
-
-
-    def set_node_liquid_water_content(self, idx, liquid_water_content):
-        """ Set liquid water content of node idx """
+    def set_node_liquid_water_content(
+        self, idx: int, liquid_water_content: float
+    ):
+        """Set liquid water content of a node."""
         self.grid[idx].set_layer_liquid_water_content(liquid_water_content)
 
-
-
-    def set_liquid_water_content(self, liquid_water_content):
-        """ Set the liquid water content profile """
+    def set_liquid_water_content(self, liquid_water_content: np.ndarray):
+        """Set the liquid water content profile."""
         for idx in range(self.number_nodes):
-            self.grid[idx].set_layer_liquid_water_content(liquid_water_content[idx])
+            self.grid[idx].set_layer_liquid_water_content(
+                liquid_water_content[idx]
+            )
 
-
-    def set_node_ice_fraction(self, idx, ice_fraction):
-        """ Set liquid ice_fraction of node idx """
+    def set_node_ice_fraction(self, idx: int, ice_fraction: float):
+        """Set liquid ice_fraction of a node."""
         self.grid[idx].set_layer_ice_fraction(ice_fraction)
 
-
-    def set_ice_fraction(self, ice_fraction):
-        """ Set the ice_fraction profile """
+    def set_ice_fraction(self, ice_fraction: np.ndarray):
+        """Set the ice fraction profile."""
         for idx in range(self.number_nodes):
             self.grid[idx].set_layer_ice_fraction(ice_fraction[idx])
 
-
-    def set_node_refreeze(self, idx, refreeze):
-        """ Set refreezing of node idx """
+    def set_node_refreeze(self, idx: int, refreeze: float):
+        """Set the refreezing of a node."""
         self.grid[idx].set_layer_refreeze(refreeze)
 
-
-
-    def set_refreeze(self, refreeze):
-        """ Set the refreezing profile """
+    def set_refreeze(self, refreeze: np.ndarray):
+        """Set the refreezing profile."""
         for idx in range(self.number_nodes):
-            self.grid[idx].set_refreeze(refreeze[idx])
-
-
+            self.grid[idx].set_layer_refreeze(refreeze[idx])
 
     def get_temperature(self):
-        """ Returns the temperature profile """
-        return [self.grid[idx].get_layer_temperature() for idx in range(self.number_nodes)]
-
-
-    def get_node_temperature(self, idx):
-        """ Returns temperature of node idx """
+        """Get the temperature profile."""
+        return [
+            self.grid[idx].get_layer_temperature()
+            for idx in range(self.number_nodes)
+        ]
+
+    def get_node_temperature(self, idx: int):
+        """Get a node's temperature."""
         return self.grid[idx].get_layer_temperature()
 
-
     def get_specific_heat(self):
-        """ Returns the specific heat (air+water+ice) profile """
-        return [self.grid[idx].get_layer_specific_heat() for idx in range(self.number_nodes)]
-
-
-    def get_node_specific_heat(self, idx):
-        """ Returns specific heat (air+water+ice) of node idx """
+        """Get the specific heat capacity profile (air+water+ice)."""
+        return [
+            self.grid[idx].get_layer_specific_heat()
+            for idx in range(self.number_nodes)
+        ]
+
+    def get_node_specific_heat(self, idx: int):
+        """Get a node's specific heat capacity (air+water+ice)."""
         return self.grid[idx].get_layer_specific_heat()
 
-
     def get_height(self):
-        """ Returns the heights of the layers """
-        return [self.grid[idx].get_layer_height() for idx in range(self.number_nodes)]
-
+        """Get the heights of all the layers."""
+        return [
+            self.grid[idx].get_layer_height()
+            for idx in range(self.number_nodes)
+        ]
 
     def get_snow_heights(self):
-        """ Returns the heights of the snow layers """
-        return [self.grid[idx].get_layer_height() for idx in range(self.get_number_snow_layers())]
-
+        """Get the heights of the snow layers."""
+        return [
+            self.grid[idx].get_layer_height()
+            for idx in range(self.get_number_snow_layers())
+        ]
 
     def get_ice_heights(self):
-        """ Returns the heights of the ice layers """
-        return [self.grid[idx].get_layer_height() for idx in range(self.number_nodes) if (self.get_node_density(idx)>=snow_ice_threshold)]
-
-
-    def get_node_height(self, idx):
-        """ Returns layer height of node idx """
+        """Get the heights of the ice layers."""
+        return [
+            self.grid[idx].get_layer_height()
+            for idx in range(self.number_nodes)
+            if (self.get_node_density(idx) >= constants.snow_ice_threshold)
+        ]
+
+    def get_node_height(self, idx: int):
+        """Get a node's layer height."""
         return self.grid[idx].get_layer_height()
 
-
-    def get_node_density(self, idx):
-        """ Returns density of node idx """
+    def get_node_density(self, idx: int):
+        """Get a node's density."""
         return self.grid[idx].get_layer_density()
 
-
     def get_density(self):
-        """ Returns the rho profile """
-        return [self.grid[idx].get_layer_density() for idx in range(self.number_nodes)]
-
-
-    def get_node_liquid_water_content(self, idx):
-        """ Returns density of node idx """
+        """Get the density profile."""
+        return [
+            self.grid[idx].get_layer_density()
+            for idx in range(self.number_nodes)
+        ]
+
+    def get_node_liquid_water_content(self, idx: int):
+        """Get a node's liquid water content."""
         return self.grid[idx].get_layer_liquid_water_content()
 
-
     def get_liquid_water_content(self):
-        """ Returns the rho profile """
-        return [self.grid[idx].get_layer_liquid_water_content() for idx in range(self.number_nodes)]
-
-
-    def get_node_ice_fraction(self, idx):
-        """ Returns ice fraction of node idx """
+        """Get a profile of the liquid water content."""
+        return [
+            self.grid[idx].get_layer_liquid_water_content()
+            for idx in range(self.number_nodes)
+        ]
+
+    def get_node_ice_fraction(self, idx: int):
+        """Get a node's ice fraction."""
         return self.grid[idx].get_layer_ice_fraction()
 
-
     def get_ice_fraction(self):
-        """ Returns the liquid water profile """
-        return [self.grid[idx].get_layer_ice_fraction() for idx in range(self.number_nodes)]
-
-
-    def get_node_irreducible_water_content(self, idx):
-        """ Returns irreducible water content of node idx """
+        """Get a profile of the ice fraction."""
+        return [
+            self.grid[idx].get_layer_ice_fraction()
+            for idx in range(self.number_nodes)
+        ]
+
+    def get_node_irreducible_water_content(self, idx: int):
+        """Get a node's irreducible water content."""
         return self.grid[idx].get_layer_irreducible_water_content()
 
-
     def get_irreducible_water_content(self):
-        """ Returns the irreducible water content profile """
-        return [self.grid[idx].get_layer_irreducible_water_content() for idx in range(self.number_nodes)]
-
-
-    def get_node_cold_content(self, idx):
-        """ Returns cold content of node idx """
+        """Get a profile of the irreducible water content."""
+        return [
+            self.grid[idx].get_layer_irreducible_water_content()
+            for idx in range(self.number_nodes)
+        ]
+
+    def get_node_cold_content(self, idx: int):
+        """Get a node's cold content."""
         return self.grid[idx].get_layer_cold_content()
 
-
     def get_cold_content(self):
-        """ Returns the cold content profile """
-        return [self.grid[idx].get_layer_cold_content() for idx in range(self.number_nodes)]
-
-
-    def get_node_porosity(self, idx):
-        """ Returns porosity of node idx """
+        """Get the cold content profile."""
+        return [
+            self.grid[idx].get_layer_cold_content()
+            for idx in range(self.number_nodes)
+        ]
+
+    def get_node_porosity(self, idx: int):
+        """Get a node's porosity."""
         return self.grid[idx].get_layer_porosity()
 
-
     def get_porosity(self):
-        """ Returns the porosity profile """
-        return [self.grid[idx].get_layer_porosity() for idx in range(self.number_nodes)]
-
-
-    def get_node_thermal_conductivity(self, idx):
-        """ Returns the thermal conductivity of node idx """
+        """Get the porosity profile."""
+        return [
+            self.grid[idx].get_layer_porosity()
+            for idx in range(self.number_nodes)
+        ]
+
+    def get_node_thermal_conductivity(self, idx: int):
+        """Get a node's thermal conductivity."""
         return self.grid[idx].get_layer_thermal_conductivity()
 
-
     def get_thermal_conductivity(self):
-        """ Returns the thermal conductivity profile """
-        return [self.grid[idx].get_layer_thermal_conductivity() for idx in range(self.number_nodes)]
-
-    def get_node_thermal_diffusivity(self, idx):
-        """ Returns the thermal diffusivity of node idx """
+        """Get the thermal conductivity profile."""
+        return [
+            self.grid[idx].get_layer_thermal_conductivity()
+            for idx in range(self.number_nodes)
+        ]
+
+    def get_node_thermal_diffusivity(self, idx: int):
+        """Get a node's thermal diffusivity."""
         return self.grid[idx].get_layer_thermal_diffusivity()
 
-
     def get_thermal_diffusivity(self):
-        """ Returns the thermal diffusivity profile """
-        return [self.grid[idx].get_layer_thermal_diffusivity() for idx in range(self.number_nodes)]
-    
-    
-    def get_node_refreeze(self, idx):
-        """ Returns refreezing of node idx """
+        """Get the thermal diffusivity profile"""
+        return [
+            self.grid[idx].get_layer_thermal_diffusivity()
+            for idx in range(self.number_nodes)
+        ]
+
+    def get_node_refreeze(self, idx: int):
+        """Get the amount of refrozen water in a node."""
         return self.grid[idx].get_layer_refreeze()
 
-
     def get_refreeze(self):
-        """ Returns the refreezing profile """
-        return [self.grid[idx].get_layer_refreeze() for idx in range(self.number_nodes)]
-
+        """Get the profile of refrozen water."""
+        return [
+            self.grid[idx].get_layer_refreeze()
+            for idx in range(self.number_nodes)
+        ]
 
-    def get_node_depth(self, idx):
+    def get_node_depth(self, idx: int):
         d = 0
-        for i in range(idx+1):
-            if i==0:
-                d = d + self.get_node_height(i)/2.0
+        for i in range(idx + 1):
+            if i == 0:
+                d = d + self.get_node_height(i) / 2.0
             else:
-                d = d + self.get_node_height(i-1)/2.0 + self.get_node_height(i)/2.0
+                d = (
+                    d
+                    + self.get_node_height(i - 1) / 2.0
+                    + self.get_node_height(i) / 2.0
+                )
         return d
 
-
     def get_depth(self):
-        """ Returns depth profile """
+        """Returns depth profile."""
         return [self.get_node_depth(idx) for idx in range(self.number_nodes)]
 
-
     def get_total_snowheight(self, verbose=False):
-        """ Get the total snowheight (density<snow_ice_threshold)"""
-        snowheights = [self.grid[idx].get_layer_height() for idx in range(self.number_nodes) if self.get_node_density(idx)<snow_ice_threshold]
-        return np.sum(np.array(snowheights))	#numba needs to be able to determine type of list contents 
- 
-    
+        """Get the total snowheight (density<snow_ice_threshold)."""
+        snowheights = [
+            self.grid[idx].get_layer_height()
+            for idx in range(self.number_nodes)
+            if self.get_node_density(idx) < constants.snow_ice_threshold
+        ]
+        return np.sum(
+            np.array(snowheights)
+        )  # numba needs to be able to determine type of list contents
+
     def get_total_height(self, verbose=False):
-        """ Get the total domain height """
+        """Get the total domain height."""
         total = [self.get_node_height(idx) for idx in range(self.number_nodes)]
         return np.sum(np.array(total))
 
-
     def get_number_snow_layers(self):
-        """ Get the number of snow layers (density<snow_ice_threshold)"""
-        nlayers = [1 for idx in range(self.number_nodes) if self.get_node_density(idx)<snow_ice_threshold]
+        """Get the number of snow layers (density<snow_ice_threshold)."""
+        nlayers = [
+            1
+            for idx in range(self.number_nodes)
+            if self.get_node_density(idx) < constants.snow_ice_threshold
+        ]
         return int(np.sum(np.array(nlayers)))
-    
 
     def get_number_layers(self):
-        """ Get the number of layers"""
-        return (self.number_nodes)
-
-
+        """Get the number of layers."""
+        return self.number_nodes
 
     def info(self):
-        """ Print some information on grid """
+        """Print some information on grid."""
 
         print("******************************")
-        print("Number of nodes:",self.number_nodes)
+        print("Number of nodes:", self.number_nodes)
         print("******************************")
 
         tmp = 0
         for i in range(self.number_nodes):
             tmp = tmp + self.get_node_height(i)
 
-        print("Grid consists of",self.number_nodes,"nodes")
-        print("Total domain depth is",tmp,"m")
-
-
+        print("Grid consists of", self.number_nodes, "nodes")
+        print("Total domain depth is", tmp, "m")
 
     def grid_info(self, n=-999):
-        """ The function prints the state of the snowpack
-            Args:
-                n   : number of nodes to plot (from top)
+        """Print the state of the snowpack.
+
+        Parameters
+        ----------
+        n : int
+            Number of nodes to plot from top.
         """
-        if (n==-999):
+        if n == -999:
             n = self.number_nodes
 
-        print("Node no., Layer height [m], Temperature [K], Density [kg m^-3], \
+        print(
+            "Node no., Layer height [m], Temperature [K], Density [kg m^-3], \
                LWC [-], LW [m], CC [J m^-2], Porosity [-], Refreezing [m w.e.], \
-	       Irreducible water content [-]")
+               Irreducible water content [-]"
+        )
 
         for i in range(n):
-            print(i, self.get_node_height(i), self.get_node_temperature(i), self.get_node_density(i), 
-                     self.get_node_liquid_water_content(i), self.get_node_cold_content(i),
-                     self.get_node_porosity(i), self.get_node_refreeze(i), self.get_node_irreducible_water_content(i))
-
-
+            print(
+                i,
+                self.get_node_height(i),
+                self.get_node_temperature(i),
+                self.get_node_density(i),
+                self.get_node_liquid_water_content(i),
+                self.get_node_cold_content(i),
+                self.get_node_porosity(i),
+                self.get_node_refreeze(i),
+                self.get_node_irreducible_water_content(i),
+            )
 
     def grid_info_screen(self, n=-999):
-        """ The function prints the state of the snowpack
-            Args:
-                n   : number of nodes to plot (from top)
+        """Prints the state of the snowpack.
+        
+        Parameters
+        ----------
+        n : int
+            Number of nodes to plot from top.
         """
-        if (n==-999):
+        if n == -999:
             n = self.number_nodes
 
-        print("Node no., Layer height [m], Temperature [K], Density [kg m^-3], LWC [-], \
-               Retention [-], CC [J m^-2], Porosity [-], Refreezing [m w.e.]")
+        print(
+            "Node no., Layer height [m], Temperature [K], Density [kg m^-3], LWC [-], \
+               Retention [-], CC [J m^-2], Porosity [-], Refreezing [m w.e.]"
+        )
 
         for i in range(n):
-            print(i, self.get_node_height(i), self.get_node_temperature(i),
-                  self.get_node_density(i), self.get_node_liquid_water_content(i), 
-                  self.get_node_irreducible_water_content(i), self.get_node_cold_content(i),
-                  self.get_node_porosity(i), self.get_node_refreeze(i))
-
-
+            print(
+                i,
+                self.get_node_height(i),
+                self.get_node_temperature(i),
+                self.get_node_density(i),
+                self.get_node_liquid_water_content(i),
+                self.get_node_irreducible_water_content(i),
+                self.get_node_cold_content(i),
+                self.get_node_porosity(i),
+                self.get_node_refreeze(i),
+            )
 
     def grid_check(self, level=1):
-        """ The function checks the grid
-            Args:
-                n   : number of nodes to plot (from top)
+        """Checks the grid.
+
+        Parameters
+        ----------
+            level : int
+                Level number.
         """
-        #if level == 1:
+        # if level == 1:
         #    self.check_layer_property(self.get_height(), 'thickness', 1.01, -0.001)
         #    self.check_layer_property(self.get_temperature(), 'temperature', 273.2, 100.0)
         #    self.check_layer_property(self.get_density(), 'density', 918, 100)
@@ -898,10 +1052,17 @@ def grid_check(self, level=1):
         #    #self.check_layer_property(self.get_cold_content(), 'CC', 1000, -10**8)
         #    self.check_layer_property(self.get_porosity(), 'Porosity', 0.8, -0.00001)
         #    self.check_layer_property(self.get_refreeze(), 'Refreezing', 0.5, 0.0)
+        pass
 
-
-
-    def check_layer_property(self, property, name, maximum, minimum, n=-999, level=1):
+    def check_layer_property(
+        self, property, name, maximum, minimum, n=-999, level=1
+    ):
         if np.nanmax(property) > maximum or np.nanmin(property) < minimum:
-            print(str.capitalize(name),'max',np.nanmax(property),'min',np.nanmin(property))
+            print(
+                str.capitalize(name),
+                "max",
+                np.nanmax(property),
+                "min",
+                np.nanmin(property),
+            )
             os._exit()
diff --git a/cosipy/cpkernel/init.py b/cosipy/cpkernel/init.py
index e6ef4926..4027d66c 100644
--- a/cosipy/cpkernel/init.py
+++ b/cosipy/cpkernel/init.py
@@ -1,11 +1,13 @@
 import sys
+
 import numpy as np
+
 from constants import *
-from config import *
-from cosipy.cpkernel.grid import *
+from cosipy.cpkernel.grid import Grid
+
 
 def init_snowpack(DATA):
-    ''' INITIALIZATION '''
+    """INITIALIZATION"""
 
     ##--------------------------------------
     ## Check for WRF data
@@ -17,7 +19,7 @@ def init_snowpack(DATA):
     else: 
         initial_snowheight = initial_snowheight_constant
     temperature_top = np.minimum(DATA.T2.values[0], 273.16)
-	
+
     #--------------------------------------
     # Do the vertical interpolation
     #--------------------------------------
@@ -41,14 +43,14 @@ def init_snowpack(DATA):
             layer_densities = np.array([initial_top_density_snowpack-drho*(initial_snowheight/nlayers)*i for i in range(1,nlayers+1)])
             layer_T = np.array([temperature_top-dT*(initial_snowheight/nlayers)*i for i in range(1,nlayers+1)])
             layer_liquid_water = np.zeros(nlayers)
-	    
+
         # add glacier
         nlayers = int(initial_glacier_height/initial_glacier_layer_heights)
         layer_heights = np.append(layer_heights, np.ones(nlayers)*initial_glacier_layer_heights)
         layer_densities = np.append(layer_densities, np.ones(nlayers)*ice_density)
         layer_T = np.append(layer_T, [layer_T[-1]-dT*initial_glacier_layer_heights*i for i in range(1,nlayers+1)])
         layer_liquid_water = np.append(layer_liquid_water, np.zeros(nlayers))
-	
+
     else:
     
         # only glacier
@@ -58,7 +60,7 @@ def init_snowpack(DATA):
         layer_densities = np.ones(nlayers)*ice_density
         layer_T = np.array([temperature_top-dT*initial_glacier_layer_heights*i for i in range(1,nlayers+1)])
         layer_liquid_water = np.zeros(nlayers)
-	
+
     # Initialize grid, the grid class contains all relevant grid information
     GRID = Grid(np.array(layer_heights, dtype=np.float64), np.array(layer_densities, dtype=np.float64), 
                 np.array(layer_T, dtype=np.float64), np.array(layer_liquid_water, dtype=np.float64),
@@ -72,7 +74,7 @@ def load_snowpack(GRID_RESTART):
     """ Initialize grid from restart file """
 
     # Number of layers
-    num_layers = np.int(GRID_RESTART.NLAYERS.values)
+    num_layers = int(GRID_RESTART.NLAYERS.values)
    
     # Init layer height
     # Weird slicing position to accommodate NestedNamespace in WRF_X_CSPY
diff --git a/cosipy/cpkernel/io.py b/cosipy/cpkernel/io.py
index 79a08e13..fccfea15 100644
--- a/cosipy/cpkernel/io.py
+++ b/cosipy/cpkernel/io.py
@@ -2,17 +2,17 @@
  This file reads the input data (model forcing) and write the output to netcdf file
 """
 
+import configparser
 import os
 import sys
-import xarray as xr
-import pandas as pd
+
 import numpy as np
-import time
-from datetime import timedelta
-from cosipy.modules.radCor import correctRadiation
+import pandas as pd
+import xarray as xr
+
+from config import *
 from constants import *
-from config import * 
-import configparser
+
 
 class IOClass:
 
@@ -42,7 +42,7 @@ def __init__(self, DATA=None):
     def create_data_file(self):
         """ Returns the DATA xarray dataset"""
     
-        if (restart==True):
+        if restart:
             print('--------------------------------------------------------------')
             print('\t RESTART FROM PREVIOUS STATE')
             print('-------------------------------------------------------------- \n')
@@ -168,7 +168,7 @@ def init_data_dataset(self):
         
         # Define a auxiliary function to check the validity of the data
         def check(field, max, min):
-            '''Check the validity of the input data '''
+            """Check the validity of the input data."""
             if np.nanmax(field) > max or np.nanmin(field) < min:
                 print('Please check the input data, its seems they are out of range %s MAX: %.2f MIN: %.2f \n' % (str.capitalize(field.name), np.nanmax(field), np.nanmin(field)))
         # Check if data is within valid bounds
@@ -426,8 +426,8 @@ def copy_local_to_global(self,y,x,local_RAIN,local_SNOWFALL,local_LWin,local_LWo
                              local_MB, local_surfMB,local_Q,local_SNOWHEIGHT,local_TOTALHEIGHT,local_TS,local_ALBEDO, \
                              local_LAYERS,local_ME,local_intMB,local_EVAPORATION,local_SUBLIMATION,local_CONDENSATION, \
                              local_DEPOSITION,local_REFREEZE,local_subM,local_Z0,local_surfM,local_MOL,local_LAYER_HEIGHT, \
-			     local_LAYER_RHO,local_LAYER_T,local_LAYER_LWC,local_LAYER_CC,local_LAYER_POROSITY, \
-			     local_LAYER_ICE_FRACTION,local_LAYER_IRREDUCIBLE_WATER,local_LAYER_REFREEZE):
+                             local_LAYER_RHO,local_LAYER_T,local_LAYER_LWC,local_LAYER_CC,local_LAYER_POROSITY, \
+                             local_LAYER_ICE_FRACTION,local_LAYER_IRREDUCIBLE_WATER,local_LAYER_REFREEZE):
 
         if ('RAIN' in self.atm):
             self.RAIN[:,y,x] = local_RAIN
@@ -564,7 +564,7 @@ def write_results_to_file(self):
             self.add_variable_along_latlontime(self.RESULT, self.surfM, 'surfM', 'm w.e.', 'Surface melt') 
         if ('MOL' in self.internal):
             self.add_variable_along_latlontime(self.RESULT, self.MOL, 'MOL', '', 'Monin Obukhov length') 
-	            
+
         if full_field:
             if ('HEIGHT' in self.full):
                 self.add_variable_along_latlonlayertime(self.RESULT, self.LAYER_HEIGHT, 'LAYER_HEIGHT', 'm', 'Layer height') 
@@ -585,6 +585,19 @@ def write_results_to_file(self):
             if ('REFREEZE' in self.full):
                 self.add_variable_along_latlonlayertime(self.RESULT, self.LAYER_REFREEZE, 'LAYER_REFREEZE', 'm w.e.', 'Refreezing') 
 
+    def create_empty_restart(self) -> xr.Dataset:
+        """Create an empty dataset for the RESTART attribute.
+
+        Returns:
+            Empty xarray dataset with coordinates from self.DATA.
+        """
+        dataset = xr.Dataset()
+        dataset.coords['time'] = self.DATA.coords['time'][-1]
+        dataset.coords['lat'] = self.DATA.coords['lat']
+        dataset.coords['lon'] = self.DATA.coords['lon']
+        dataset.coords['layer'] = np.arange(max_layers)
+
+        return dataset
 
     #----------------------------------------------
     # Initializes the restart xarray dataset
@@ -596,13 +609,9 @@ def init_restart_dataset(self):
             
             self.RESTART  ::  xarray structure"""
         
-        self.RESTART = xr.Dataset()
-        self.RESTART.coords['time'] = self.DATA.coords['time'][-1]
-        self.RESTART.coords['lat'] = self.DATA.coords['lat']
-        self.RESTART.coords['lon'] = self.DATA.coords['lon']
-        self.RESTART.coords['layer'] = np.arange(max_layers)
+        self.RESTART = self.create_empty_restart()
     
-        print('Restart ddataset ... ok \n')
+        print('Restart dataset ... ok \n')
         print('--------------------------------------------------------------\n')
         
         return self.RESTART
@@ -638,11 +647,7 @@ def create_local_restart_dataset(self):
             
             self.RESTART  ::  one-dimensional self.RESULT structure"""
     
-        self.RESTART = xr.Dataset()
-        self.RESTART.coords['time'] = self.DATA.coords['time'][-1]
-        self.RESTART.coords['lat'] = self.DATA.coords['lat']
-        self.RESTART.coords['lon'] = self.DATA.coords['lon']
-        self.RESTART.coords['layer'] = np.arange(max_layers)
+        self.RESTART = self.create_empty_restart()
         
         self.add_variable_along_scalar(self.RESTART, np.full((1), np.nan), 'NLAYERS', '-', 'Number of layers')
         self.add_variable_along_scalar(self.RESTART, np.full((1), np.nan), 'NEWSNOWHEIGHT', 'm .w.e', 'New snow height')
diff --git a/cosipy/cpkernel/node.py b/cosipy/cpkernel/node.py
index 88a6abce..19391ed5 100644
--- a/cosipy/cpkernel/node.py
+++ b/cosipy/cpkernel/node.py
@@ -1,264 +1,283 @@
-from constants import *
-import numpy as np
-
 from collections import OrderedDict
-from numba import float64  
+
+import numpy as np
+from numba import float64
 from numba.experimental import jitclass
 
+import constants
+
 spec = OrderedDict()
-spec['height'] = float64              
-spec['temperature'] = float64     
-spec['liquid_water_content'] = float64     
-spec['ice_fraction'] = float64
-spec['refreeze'] = float64
+spec["height"] = float64
+spec["temperature"] = float64
+spec["liquid_water_content"] = float64
+spec["ice_fraction"] = float64
+spec["refreeze"] = float64
 
 @jitclass(spec)
 class Node:
-    """ The Node-class stores the state variables of layer. 
+    """A `Node` class stores a layer's state variables.
     
-    The numerical grid consists of a list of nodes that store the information 
-    of individual layers. The class provides various setter/getter functions
-    to read or overwrite the state of individual layers. 
+    The numerical grid consists of a list of nodes which store the
+    information of individual layers. The class provides various
+    setter/getter functions to read or overwrite the state of these
+    individual layers. 
 
-    Parameters
+    Attributes
     ----------
-        height : float
-            Height of the layer [:math:`m`]
-        density : float
-            snow density [:math:`kg~m^{-3}`]
-        temperature : float
-            temperature of the layer [:math:`K`]
-        liquid_water : float
-            liquid water [:math:`m~w.e.`]
+    height : float
+        Layer height [:math:`m`].
+    snow_density : float
+        Layer snow density [:math:`kg~m^{-3}`].
+    temperature: float
+        Layer temperature [:math:`K`].
+    liquid_water_content : float
+        Liquid water content [:math:`m~w.e.`].
+    ice_fraction : float
+        Volumetric ice fraction [-].
+    refreeze : float
+        Amount of refrozen liquid water [:math:`m~w.e.`].
 
     Returns
     -------
-        Node : :py:class:`cosipy.cpkernel.node` object
-
+    Node : :py:class:`cosipy.cpkernel.node` object.
     """
 
-    def __init__(self, height, snow_density, temperature, liquid_water_content, ice_fraction=None):
-
-        # Initialize state variables 
+    def __init__(
+        self,
+        height: float,
+        snow_density: float,
+        temperature: float,
+        liquid_water_content: float,
+        ice_fraction: float = None,
+    ):
+        # Initialises state variables.
         self.height = height
         self.temperature = temperature
         self.liquid_water_content = liquid_water_content
-        
+
         if ice_fraction is None:
             # Remove weight of air from density
-            a = snow_density - (1-(snow_density/ice_density))*air_density
-            self.ice_fraction = a/ice_density
+            a = snow_density - (1 - (snow_density / constants.ice_density)) * constants.air_density
+            self.ice_fraction = a / constants.ice_density
         else:
             self.ice_fraction = ice_fraction
 
-        self.refreeze = 0.0 
+        self.refreeze = 0.0
 
 
-    ''' GETTER FUNCTIONS '''
-    
-    #------------------------------------------
+    """GETTER FUNCTIONS"""
+
+    # -----------------------------------------
     # Getter-functions for state variables
-    #------------------------------------------
-    def get_layer_height(self):
-        """ Returns the layer height of the node.
+    # -----------------------------------------
+    def get_layer_height(self) -> float:
+        """Gets the node's layer height.
         
         Returns
         -------
-            height : float
-                Snow layer height [:math:`m`]
+        height : float
+            Snow layer height [:math:`m`].
         """
         return self.height
 
-    def get_layer_temperature(self):
-        """ Returns the snow layer temperature of the node. 
+    def get_layer_temperature(self) -> float:
+        """Gets the node's snow layer temperature.
         
         Returns
         -------
-            T : float
-                Snow layer temperature [:math:`K`]
+        T : float
+            Snow layer temperature [:math:`K`].
         """
         return self.temperature
-    
-    def get_layer_ice_fraction(self):
-        """ Returns the volumetric ice fraction of the node. 
+
+    def get_layer_ice_fraction(self) -> float:
+        """Gets the node's volumetric ice fraction.
         
         Returns
         -------
-            ice_fraction : float
-                The volumetric ice fraction [-] 
+        ice_fraction : float
+            The volumetric ice fraction [-].
         """
-        return self.ice_fraction 
-    
-    def get_layer_refreeze(self):
-        """ Returns the amount of refreezing of the node. 
+        return self.ice_fraction
+
+    def get_layer_refreeze(self) -> float:
+        """Gets the amount of refrozen water in the node.
         
         Returns
         -------
-            refreeze : float
-                Amount of water that has refreezed per time step [:math:`m~w.e.`]
+        refreeze : float
+            Amount of refrozen water per time step [:math:`m~w.e.`].
         """
         return self.refreeze
 
-
-    #----------------------------------------------
+    # ---------------------------------------------
     # Getter-functions for derived state variables
-    #----------------------------------------------
-    def get_layer_density(self):
-        """ Returns the mean density including ice and liquid of the node. 
+    # ---------------------------------------------
+    def get_layer_density(self) -> float:
+        """Gets the node's mean density including ice and liquid.
 
         Returns
         -------
-            rho : float
-                Snow density [:math:`kg~m^{-3}`]
+        rho : float
+            Snow density [:math:`kg~m^{-3}`].
         """
-        return self.get_layer_ice_fraction()*ice_density + self.get_layer_liquid_water_content()*water_density + self.get_layer_air_porosity()*air_density
-    
-    def get_layer_air_porosity(self):
-        """ Returnis the ice fraction of the node.
+        return (
+            self.get_layer_ice_fraction() * constants.ice_density
+            + self.get_layer_liquid_water_content() * constants.water_density
+            + self.get_layer_air_porosity() * constants.air_density
+        )
+
+    def get_layer_air_porosity(self) -> float:
+        """Gets the fraction of air in the node.
 
         Returns
         -------
-            porosity : float
-                Air porosity [:math:`m`]
+        porosity : float
+            Air porosity [:math:`m`].
         """
         return max(0.0, 1 - self.get_layer_liquid_water_content() - self.get_layer_ice_fraction())
-    
-    def get_layer_specific_heat(self):
-        """ Returns the volumetric averaged specific heat of the node. 
+
+    def get_layer_specific_heat(self) -> float:
+        """Gets the node's volumetric averaged specific heat capacity.
 
         Returns
         -------
-            cp : float
-                Specific heat [:math:`J~kg^{-1}~K^{-1}`]
+        cp : float
+            Specific heat capacity [:math:`J~kg^{-1}~K^{-1}`].
         """
-        return self.get_layer_ice_fraction()*spec_heat_ice + self.get_layer_air_porosity()*spec_heat_air + self.get_layer_liquid_water_content()*spec_heat_water
+        return self.get_layer_ice_fraction()*constants.spec_heat_ice + self.get_layer_air_porosity()*constants.spec_heat_air + self.get_layer_liquid_water_content()*constants.spec_heat_water
 
-    def get_layer_liquid_water_content(self):
-        """ Returns the liquid water content of the node.
+    def get_layer_liquid_water_content(self) -> float:
+        """Gets the node's liquid water content.
 
         Returns
         -------
-            lwc : float
-                Liquid water content [-]
+        lwc : float
+            Liquid water content [-].
         """
-        return self.liquid_water_content 
-    
-    def get_layer_irreducible_water_content(self):
-        """ Returns the irreducible water content of the node. 
+        return self.liquid_water_content
+
+    def get_layer_irreducible_water_content(self) -> float:
+        """Gets the node's irreducible water content.
 
         Returns
         -------
-            ret : float
-                Irreducible water content [-]
+        theta_e : float
+            Irreducible water content [-].
         """
-        if (self.get_layer_ice_fraction() <= 0.23):
-            theta_e = 0.0264 + 0.0099*((1-self.get_layer_ice_fraction())/self.get_layer_ice_fraction()) 
-        elif (self.get_layer_ice_fraction() > 0.23) & (self.get_layer_ice_fraction() <= 0.812):
-            theta_e = 0.08 - 0.1023*(self.get_layer_ice_fraction()-0.03)
+
+        ice_fraction = self.get_layer_ice_fraction()
+        if ice_fraction <= 0.23:
+            theta_e = 0.0264 + 0.0099 * ((1 - ice_fraction) / ice_fraction)
+        elif (ice_fraction > 0.23) & (ice_fraction <= 0.812):
+            theta_e = 0.08 - 0.1023 * (ice_fraction - 0.03)
         else:
             theta_e = 0.0
-        return theta_e 
-    
-    def get_layer_cold_content(self):
-        """ Returns the cold content of the node. 
+        return theta_e
+
+    def get_layer_cold_content(self) -> float:
+        """Gets the node's cold content.
 
         Returns
         -------
-            cc : float
-                Cold content [:math:`J~m^{-2}`]
+        cc : float
+            Cold content [:math:`J~m^{-2}`].
         """
-        return -self.get_layer_specific_heat() * self.get_layer_density() * self.get_layer_height() * (self.get_layer_temperature()-zero_temperature)
-    
-    def get_layer_porosity(self):
-        """ Returns the porosity of the node. 
+        return -self.get_layer_specific_heat() * self.get_layer_density() * self.get_layer_height() * (self.get_layer_temperature()-constants.zero_temperature)
+
+    def get_layer_porosity(self) -> float:
+        """Gets the node's porosity.
 
         Returns
         -------
-            por : float
-                Air porosity [-]
+        porosity : float
+            Air porosity [-].
         """
         return 1-self.get_layer_ice_fraction()-self.get_layer_liquid_water_content()
-   
-    def get_layer_thermal_conductivity(self):
-        """ Returns the volumetric weighted thermal conductivity of the node.
+
+    def get_layer_thermal_conductivity(self) -> float:
+        """Gets the node's volumetric weighted thermal conductivity.
 
         Returns
         -------
-            lambda : float
-                Thermal conductivity [:math:`W~m^{-1}~K^{-1}`]
+        lam : float
+            Thermal conductivity [:math:`W~m^{-1}~K^{-1}`].
         """
         methods_allowed = ['bulk', 'empirical']
-        if thermal_conductivity_method == 'bulk':
-            lam = self.get_layer_ice_fraction()*k_i + self.get_layer_air_porosity()*k_a + self.get_layer_liquid_water_content()*k_w
-        elif thermal_conductivity_method == 'empirical':
+        if constants.thermal_conductivity_method == 'bulk':
+            lam = self.get_layer_ice_fraction()*constants.k_i + self.get_layer_air_porosity()*constants.k_a + self.get_layer_liquid_water_content()*constants.k_w
+        elif constants.thermal_conductivity_method == 'empirical':
             lam = 0.021 + 2.5 * np.power((self.get_layer_density()/1000),2)
         else:
-            raise ValueError("Thermal conductivity method = \"{:s}\" is not allowed, must be one of {:s}".format(thermal_conductivity_method, ", ".join(methods_allowed)))
+            message = ("Thermal conductivity method =",
+                       f"{constants.thermal_conductivity_method}",
+                       f"is not allowed, must be one of",
+                       f"{', '.join(methods_allowed)}")
+            raise ValueError(" ".join(message))
         return lam
 
-    def get_layer_thermal_diffusivity(self):
-        """ Returns the thermal diffusivity of the node. 
+    def get_layer_thermal_diffusivity(self) -> float:
+        """Gets the node's thermal diffusivity.
 
         Returns
         -------
-            K : float
-                Thermal diffusivity [:math:`m^{2}~s^{-1}`]
+        K : float
+            Thermal diffusivity [:math:`m^{2}~s^{-1}`].
         """
         K = self.get_layer_thermal_conductivity()/(self.get_layer_density()*self.get_layer_specific_heat())
         return K
 
+    """SETTER FUNCTIONS"""
 
-    ''' SETTER FUNCTIONS '''
-
-    #----------------------------------------------
+    # ---------------------------------------------
     # Setter-functions for derived state variables
-    #----------------------------------------------
-    def set_layer_height(self, height):
-        """ Sets the layer height of the node. 
+    # ---------------------------------------------
+    def set_layer_height(self, height: float):
+        """Sets the node's layer height.
         
         Parameters
         ----------
-            height : float
-                Layer height [:math:`m`]
+        height : float
+            Layer height [:math:`m`].
         """
         self.height = height
 
-    def set_layer_temperature(self, T):
-        """ Sets the mean temperature of the node.
+    def set_layer_temperature(self, T: float):
+        """Sets the node's mean temperature.
 
         Parameters
         ----------
-            T : float
-                Layer temperature [:math:`K`]
+        T : float
+            Layer temperature [:math:`K`].
         """
         self.temperature = T
 
-    def set_layer_liquid_water_content(self, lwc):
-        """ Sets the liquid water content of the node.
+    def set_layer_liquid_water_content(self, lwc: float):
+        """Sets the node's liquid water content.
 
         Parameters
         ----------
-            lwc : float
-                Liquid water content [-]
+        lwc : float
+            Liquid water content [-].
         """
         self.liquid_water_content = lwc
-    
-    def set_layer_ice_fraction(self, ifr):
-        """ Sets the ice fraction of the node. 
+
+    def set_layer_ice_fraction(self, ifr: float):
+        """Sets the node's volumetric ice fraction.
 
         Parameters
         ----------
-            ifr : float
-                Ice fraction [-]
+        ifr : float
+            Volumetric ice fraction [-].
         """
         self.ice_fraction = ifr
-    
-    def set_layer_refreeze(self, refr):
-        """ Sets the amount of water refreezed of the node.
+
+    def set_layer_refreeze(self, refr: float):
+        """Sets the amount of refrozen water in the node.
 
         Parameters
         ----------
-            refr : float
-                Amount of refreezed water [:math:`m~w.e.`]
+        refr : float
+            Amount of refrozen water [:math:`m~w.e.`].
         """
         self.refreeze = refr
diff --git a/cosipy/modules/albedo.py b/cosipy/modules/albedo.py
index 22e12263..0281eb4a 100644
--- a/cosipy/modules/albedo.py
+++ b/cosipy/modules/albedo.py
@@ -1,65 +1,75 @@
 import numpy as np
-from constants import albedo_method, albedo_fresh_snow, albedo_firn, albedo_ice, \
-                      albedo_mod_snow_aging, albedo_mod_snow_depth, snow_ice_threshold, zero_temperature, dt, t_star_dry, t_star_wet, t_star_K, t_star_cutoff
 
-def updateAlbedo(GRID,surface_temperature,albedo_snow):
-    """ This methods updates the albedo """
-    albedo_allowed = ['Oerlemans98','Bougamont05']
-    if albedo_method == 'Oerlemans98':
+import constants
+
+
+def updateAlbedo(GRID, surface_temperature, albedo_snow):
+    """Updates the albedo."""
+    albedo_allowed = ["Oerlemans98", "Bougamont05"]
+    if constants.albedo_method == "Oerlemans98":
         alphaMod = method_Oerlemans(GRID)
-    elif albedo_method == 'Bougamont05':
-        alphaMod, albedo_snow = method_Bougamont(GRID,surface_temperature,albedo_snow)
+    elif constants.albedo_method == "Bougamont05":
+        alphaMod, albedo_snow = method_Bougamont(GRID, surface_temperature, albedo_snow)
     else:
-        raise ValueError("Albedo method = \"{:s}\" is not allowed, must be one of {:s}".format(albedo_method, ", ".join(albedo_allowed)))
+        error_message = (
+            f'Albedo method = "{constants.albedo_method}"',
+            f"is not allowed, must be one of",
+            f'{", ".join(albedo_allowed)}',
+        )
+        raise ValueError(" ".join(error_message))
 
     return alphaMod, albedo_snow
 
 
 def method_Oerlemans(GRID):
-
     # Get hours since the last snowfall
     # First get fresh snow properties (height and timestamp)
-    fresh_snow_height, fresh_snow_timestamp, _  = GRID.get_fresh_snow_props()
-    
+    fresh_snow_height, fresh_snow_timestamp, _ = GRID.get_fresh_snow_props()
+
     # Get time difference between last snowfall and now
-    hours_since_snowfall = (fresh_snow_timestamp)/3600.0
+    hours_since_snowfall = (fresh_snow_timestamp) / 3600.0
 
     # If fresh snow disappears faster than the snow ageing scale then set the hours_since_snowfall
     # to the old values of the underlying snowpack
-    if (hours_since_snowfall<(albedo_mod_snow_aging*24)) & (fresh_snow_height<0.0):
+    if (hours_since_snowfall < (constants.albedo_mod_snow_aging * 24)) & (
+        fresh_snow_height <= 0.0
+    ):
         GRID.set_fresh_snow_props_to_old_props()
-        fresh_snow_height, fresh_snow_timestamp, _  = GRID.get_fresh_snow_props()
-        
+        fresh_snow_height, fresh_snow_timestamp, _ = GRID.get_fresh_snow_props()
+
         # Update time difference between last snowfall and now
-        hours_since_snowfall = (fresh_snow_timestamp)/3600.0
+        hours_since_snowfall = (fresh_snow_timestamp) / 3600.0
 
     # Check if snow or ice
-    if (GRID.get_node_density(0) <= snow_ice_threshold):
-        
+    if GRID.get_node_density(0) <= constants.snow_ice_threshold:
         # Get current snowheight from layer height
-        h = GRID.get_total_snowheight() #np.sum(GRID.get_height()[0:idx])
+        h = GRID.get_total_snowheight()  # np.sum(GRID.get_height()[0:idx])
 
         # Surface albedo according to Oerlemans & Knap 1998, JGR)
-        alphaSnow = albedo_firn + (albedo_fresh_snow - albedo_firn) *  np.exp((-hours_since_snowfall) / (albedo_mod_snow_aging * 24.0))
-        alphaMod = alphaSnow + (albedo_ice - alphaSnow) *  np.exp((-1.0*h) / (albedo_mod_snow_depth / 100.0))
+        alphaSnow = constants.albedo_firn + (
+            constants.albedo_fresh_snow - constants.albedo_firn
+        ) * np.exp((-hours_since_snowfall) / (constants.albedo_mod_snow_aging * 24.0))
+        alphaMod = alphaSnow + (constants.albedo_ice - alphaSnow) * np.exp(
+            (-1.0 * h) / (constants.albedo_mod_snow_depth / 100.0)
+        )
 
     else:
         # If no snow cover than set albedo to ice albedo
-        alphaMod = albedo_ice
+        alphaMod = constants.albedo_ice
 
     return alphaMod
 
-def method_Bougamont(GRID,surface_temperature,albedo_snow):
 
+def method_Bougamont(GRID, surface_temperature, albedo_snow):
     # Get hours since the last snowfall
     # First get fresh snow properties (height and timestamp)
-    _ , fresh_snow_timestamp, _  = GRID.get_fresh_snow_props()
+    _, fresh_snow_timestamp, _ = GRID.get_fresh_snow_props()
 
     # Get time difference between last snowfall and now:
     hours_since_snowfall = (fresh_snow_timestamp) / 3600.0
 
     # Convert integration time from seconds to days:
-    dt_days = dt / 86400.0
+    dt_days = constants.dt / 86400.0
 
     # Note: accounting for disapearance of uppermost fresh snow layer difficult due to non-constant decay rate. Unsure how to implement.
 
@@ -67,37 +77,49 @@ def method_Bougamont(GRID,surface_temperature,albedo_snow):
     h = GRID.get_total_snowheight()
 
     # Check if snow or ice:
-    if (GRID.get_node_density(0) <= snow_ice_threshold):
-
-        if surface_temperature >= zero_temperature:
+    if GRID.get_node_density(0) <= constants.snow_ice_threshold:
+        if surface_temperature >= constants.zero_temperature:
             # Snow albedo decay timescale (t*) on a melting snow surface:
-            t_star = t_star_wet
+            t_star = constants.t_star_wet
         else:
             # Snow albedo decay timescale (t*) on a dry snow surface:
-            if surface_temperature < t_star_cutoff:
-                t_star = t_star_dry + (zero_temperature - t_star_cutoff) * t_star_K
+            if surface_temperature < constants.t_star_cutoff:
+                t_star = (
+                    constants.t_star_dry
+                    + (constants.zero_temperature - constants.t_star_cutoff)
+                    * constants.t_star_K
+                )
             else:
-                t_star = t_star_dry + (zero_temperature - surface_temperature) * t_star_K
+                t_star = (
+                    constants.t_star_dry
+                    + (constants.zero_temperature - surface_temperature)
+                    * constants.t_star_K
+                )
 
         # Effect of snow albedo decay due to the temporal metamorphosis of snow (Bougamont et al. 2005 - based off Oerlemans & Knap 1998):
         # Exponential function discretised in order to account for variable surface temperature-dependant decay timescales.
-        albedo_snow = albedo_snow - (albedo_snow - albedo_firn) / t_star * dt_days 
+        albedo_snow = (
+            albedo_snow - (albedo_snow - constants.albedo_firn) / t_star * dt_days
+        )
 
         # Reset if snowfall in current timestep
-        if hours_since_snowfall == 0:    
-            albedo_snow = albedo_fresh_snow
+        if hours_since_snowfall == 0:
+            albedo_snow = constants.albedo_fresh_snow
 
         # Effect of surface albedo decay due to the snow depth (Oerlemans & Knap 1998):
-        alphaMod = albedo_snow + (albedo_ice - albedo_snow) *  np.exp((-1.0 * h) / (albedo_mod_snow_depth / 100.0))
-        
+        alphaMod = albedo_snow + (constants.albedo_ice - albedo_snow) * np.exp(
+            (-1.0 * h) / (constants.albedo_mod_snow_depth / 100.0)
+        )
+
     else:
         # If no snow cover than set albedo to ice albedo
-        alphaMod = albedo_ice
+        alphaMod = constants.albedo_ice
 
     # Ensure output value is of the float data type.
     alphaMod = float(alphaMod)
 
     return alphaMod, albedo_snow
 
+
 ### idea; albedo decay like (Brock et al. 2000)? or?
 ### Schmidt et al 2017 >doi:10.5194/tc-2017-67, 2017 use the same albedo parameterisation from Oerlemans and Knap 1998 with a slight updated implementation of considering the surface temperature?
diff --git a/cosipy/modules/densification.py b/cosipy/modules/densification.py
index fa7e16a1..53435298 100644
--- a/cosipy/modules/densification.py
+++ b/cosipy/modules/densification.py
@@ -4,10 +4,12 @@
 from numba import njit
 
 def densification(GRID,SLOPE,dt):
-    """ Densification of the snowpack
+    """Densification of the snowpack.
+
     Args:
-        GRID    ::  GRID-Structure
-	dt      ::  integration time
+        GRID    ::  GRID-Structure.
+        SLOPE   ::  Slope of the surface [degrees].
+        dt      ::  integration time [s].
     """
 
     densification_allowed = ['Boone', 'Vionnet', 'empirical', 'constant']
@@ -20,7 +22,33 @@ def densification(GRID,SLOPE,dt):
     elif densification_method == 'constant':
         pass
     else:
-        raise ValueError("Densification method = \"{:s}\" is not allowed, must be one of {:s}".format(densification_method, ", ".join(densification_allowed)))
+        error_msg = (
+            f'Densification method = "{densification_method}"',
+            "is not allowed, must be one of",
+            f'{", ".join(densification_allowed)}',
+        )
+        raise ValueError(" ".join(error_msg))
+    
+@njit
+def copy_layer_profiles(GRID) -> tuple:
+    """Get a copy of the layer profiles.
+
+    `np.array` returns a copy by default and is 2x faster than np.copy
+    (which is not supported by numba).
+
+    Args:
+        GRID (Grid): Grid object.
+    Returns:
+        Profiles for height, density, temperature, liquid water content,
+        and ice fraction.
+    """
+    heights = np.array(GRID.get_height())
+    densities = np.array(GRID.get_density())
+    temperatures = np.array(GRID.get_temperature())
+    lwcs = np.array(GRID.get_liquid_water_content())
+    ice_fractions = np.array(GRID.get_ice_fraction())
+
+    return heights, densities, temperatures, lwcs, ice_fractions
 
 @njit
 def method_Boone(GRID,SLOPE,dt):
@@ -41,12 +69,7 @@ def method_Boone(GRID,SLOPE,dt):
     M_s = 0.0
 
     # Get copy of layer heights and layer densities
-    #np.array returns a copy by default and is 2x faster than np.copy (not supported by numba)
-    rho = np.array(GRID.get_density())
-    height = np.array(GRID.get_height())
-    lwc = np.array(GRID.get_liquid_water_content())
-    t = np.array(GRID.get_temperature())
-    icf = np.array(GRID.get_ice_fraction())
+    height, rho, t, lwc, icf = copy_layer_profiles(GRID)
 
     # Loop over all internal snow nodes
     for idxNode in range(0,GRID.get_number_snow_layers() , 1):
@@ -108,11 +131,7 @@ def method_Vionnet(GRID,SLOPE,dt):
     sigma = 0.0
 
     # Get copy of layer heights and layer densities
-    rho = np.array(GRID.get_density())
-    height = np.array(GRID.get_height())
-    lwc = np.array(GRID.get_liquid_water_content())
-    t = np.array(GRID.get_temperature())
-    icf = np.array(GRID.get_ice_fraction())
+    height, rho, t, lwc, icf = copy_layer_profiles(GRID)
 
     # Loop over all internal snow nodes
     for idxNode in range(0,GRID.get_number_snow_layers() , 1):
@@ -165,7 +184,7 @@ def method_Vionnet(GRID,SLOPE,dt):
 
 
 def method_empirical(GRID,SLOPE,dt):
-    """ Simple empricial snow compaction parametrization using a constant time scale. """
+    """Simple empirical snow compaction parametrization using a constant time scale."""
 
     rho_max = 600.0     # maximum attainable density [kg m^-3]
     #tau = 3.6e5         # empirical compaction time scale [s]
diff --git a/cosipy/modules/evaluation.py b/cosipy/modules/evaluation.py
index b641d388..8ea31bb6 100644
--- a/cosipy/modules/evaluation.py
+++ b/cosipy/modules/evaluation.py
@@ -1,4 +1,3 @@
-import numpy as np
 from config import eval_method, obs_type
 
 
@@ -17,6 +16,9 @@ def evaluate(stake_names, stake_data, df_):
 def rmse(stake_names, stake_data, df_):
     if (obs_type=='mb'):
         rmse = ((stake_data[stake_names].subtract(df_['mb'],axis=0))**2).mean()**.5
-    if (obs_type=='snowheight'):
+    elif (obs_type=='snowheight'):
         rmse = ((stake_data[stake_names].subtract(df_['snowheight'],axis=0))**2).mean()**.5
+    else:
+        msg = f'RMSE not implemented for obs_type="{obs_type}" in config.py.'
+        raise NotImplementedError(msg)
     return rmse
diff --git a/cosipy/modules/penetratingRadiation.py b/cosipy/modules/penetratingRadiation.py
index 6d989519..336df166 100644
--- a/cosipy/modules/penetratingRadiation.py
+++ b/cosipy/modules/penetratingRadiation.py
@@ -49,7 +49,10 @@ def method_Bintanja(GRID, SWnet, dt):
             # Compute conversion factor A
             A = (spec_heat_ice*ice_density)/(water_density*lat_heat_melting)
 
-            # Changes in volumetric contents; dtheta_w change in water fraciont and dtheate_i change in ice fraction
+            """Changes in volumetric contents:
+            * dtheta_w: change in water fraction
+            * dtheta_i: change in ice fraction
+            """
             dtheta_w = A * dT * GRID.get_node_ice_fraction(idxNode)
             dtheta_i = (water_density/ice_density) * -dtheta_w
             dh = -dtheta_i/GRID.get_node_ice_fraction(idxNode)
diff --git a/cosipy/modules/percolation.py b/cosipy/modules/percolation.py
index b4fd1306..e3296228 100644
--- a/cosipy/modules/percolation.py
+++ b/cosipy/modules/percolation.py
@@ -1,16 +1,20 @@
 import numpy as np
 from numba import njit
 
+
 @njit
-def percolation(GRID, water, dt):
-    """ Percolation and refreezing of melt water through the snow- and firn pack
+def percolation(GRID, water: float, dt: int) -> float:
+    """Percolate melt water through the snow- and firn pack.
 
-    Args:
+    Bucket method (Bartelt & Lehning, 2002).
 
-        GRID    ::  GRID-Structure 
-        water   ::  Melt water (m w.e.q.) at the surface
-        dt      ::  Integration time
+    Args:
+        GRID (Grid): Glacier data mesh.
+        water: Melt water at the surface, [m w.e.q.].
+        dt: Integration time.
 
+    Returns:
+        Percolated meltwater.
     """
 
     # convert m to mm = kg/m2, not needed because of change to fraction
@@ -18,53 +22,83 @@ def percolation(GRID, water, dt):
 
     # convert kg/m2 to kg/m3
     water = water / GRID.get_node_height(0)
-
-
     # kg/m3 to fraction
     # water = water / 1000
 
-    # initial runoff [m w.e.]
-    Q = 0  
-    
     # set liquid water of top layer (idx, LWCnew) in m
-    GRID.set_node_liquid_water_content(0, GRID.get_node_liquid_water_content(0)+float(water))
-
-    # for consistency check
-    #numba expect numpy type in np.sum()
-    total_start = np.sum(np.array(GRID.get_liquid_water_content()))
+    GRID.set_node_liquid_water_content(
+        0, GRID.get_node_liquid_water_content(0) + float(water)
+    )
 
-    
-    # Loop over all internal grid points for percolation 
-    for idxNode in range(0, GRID.number_nodes-1, 1):
-
-        # Get irreducible water content [-]
+    # Loop over all internal grid points for percolation
+    for idxNode in range(0, GRID.number_nodes - 1, 1):
         theta_e = GRID.get_node_irreducible_water_content(idxNode)
-        
-        # Get liquid water content [-]
         theta_w = GRID.get_node_liquid_water_content(idxNode)
-    
+
         # Residual volume fraction of water (m^3 which is equal to m)
         residual = np.maximum((theta_w - theta_e), 0.0)
 
-        if residual > 0:
-            # than percolate to the next layer (add to the next layer)
+        if residual > 0.0:
             GRID.set_node_liquid_water_content(idxNode, theta_e)
 
-            ### old
-            #GRID.set_node_liquid_water_content(idxNode+1, GRID.get_node_liquid_water_content(idxNode+1)+residual)
-
-            ### new: if water is pushed to next layer, because of fractions the layer heights have to be considered
+            """
+            old:
+            GRID.set_node_liquid_water_content(
+                idxNode + 1,
+                GRID.get_node_liquid_water_content(idxNode + 1) + residual,
+            )
+
+            new:
+            If water is pushed to next layer, the layer heights have to
+            be considered because of fractions.
+            """
             residual = residual * GRID.get_node_height(idxNode)
-            GRID.set_node_liquid_water_content(idxNode + 1, GRID.get_node_liquid_water_content(idxNode + 1) + residual / GRID.get_node_height(idxNode+1))
-        else: 
-            GRID.set_node_liquid_water_content(idxNode, theta_w)
+            GRID.set_node_liquid_water_content(
+                idxNode + 1,
+                GRID.get_node_liquid_water_content(idxNode + 1)
+                + residual / GRID.get_node_height(idxNode + 1),
+            )
+
+    """Runoff is equal to LWC in the last node & must be converted
+    from kg/m3 to kg/m2. Converting from fraction to kg/m3 (*1000) and
+    from mm to m (/1000) is unnecessary."""
+    Q = GRID.get_node_liquid_water_content(
+        GRID.number_nodes - 1
+    ) * GRID.get_node_height(GRID.number_nodes - 1)
+    GRID.set_node_liquid_water_content(GRID.number_nodes - 1, 0.0)
 
-    # Runoff is equal to the LWC in the last node and has to be converted from kg/m3 to kg/m2
-    # convert from fraction to kg/m3 (*1000) and from mm to m (/1000) not needed
-    Q = GRID.get_node_liquid_water_content(GRID.number_nodes-1) * GRID.get_node_height(GRID.number_nodes-1)
-    GRID.set_node_liquid_water_content(GRID.number_nodes-1, 0.0)
+    return Q
 
-    # for consistency check
-    total_end = np.sum(np.array(GRID.get_liquid_water_content()))
 
-    return Q
+@njit
+def check_lwc_conservation(GRID, start_lwc: float, runoff: float, dt: float):
+    """Check total liquid water content is conserved.
+
+    Args:
+        GRID (Grid): Glacier data mesh.
+        start_lwc: Initial total liquid water content.
+        runoff: Meltwater runoff from the lowest node.
+        dt: Integration time [s].
+    """
+
+    end_lwc = (
+        np.nansum(
+            np.array(GRID.get_liquid_water_content())
+            * np.array(GRID.get_height())
+        )
+        + runoff
+    )
+
+    if not np.isclose(start_lwc, end_lwc):
+        if (
+            GRID.new_snow_timestamp == 0.0
+        ):  # can't index xarrays directly with njit
+            snow_time = GRID.old_snow_timestamp
+        else:
+            snow_time = GRID.new_snow_timestamp
+        timestep = snow_time / dt
+        delta = start_lwc - end_lwc
+        warn_sanity = "\nWARNING: When percolating, the initial LWC is not equal to final LWC"
+        # numba doesn't support warnings, and we don't want to raise an error
+        print(f"{warn_sanity} at timestep {int(timestep)}. dLWC:")
+        print(delta)
diff --git a/cosipy/modules/radCor.py b/cosipy/modules/radCor.py
index 297abeba..6be1d6ea 100644
--- a/cosipy/modules/radCor.py
+++ b/cosipy/modules/radCor.py
@@ -1,8 +1,10 @@
 import math
-import numpy as np
+
 import metpy.calc
+import numpy as np
 from metpy.units import units
 
+
 # The following functions are needed for the radiation method Wohlfahrt 2016
 def solarFParallel(lat, lon, timezone_lon, day, hour):
     """ Calculate solar elevation, zenith and azimuth angles
@@ -13,7 +15,7 @@ def solarFParallel(lat, lon, timezone_lon, day, hour):
         lon             ::  longitude (decimal degree)
         timezone_lon    ::  longitude of standard meridian (decimal degree)
         doy             ::  day of the year (1-366)
-        hour            ::  hour of the day (decimal, e.g. 12:30 = 12.5
+        hour            ::  hour of the day (decimal, e.g. 12:30 = 12.5)
         
     Outputs:
         
@@ -161,7 +163,7 @@ def solpars(lat):
     """
 
     timecorr = np.zeros((366, 4))
-    solpars = np.zeros((366, 7))
+    solparam = np.zeros((366, 7))
 
     for j in np.arange(0, 365):
         # Time correction
@@ -175,19 +177,19 @@ def solpars(lat):
 
         # Solar parameters
         tau = 2 * math.pi * (j) / 365
-        solpars[j, 0] = tau
-        solpars[j, 1] = tau * 180 / math.pi
-        solpars[j, 2] = 1.00011 + 0.034221 * math.cos(tau) + 0.00128 * math.sin(tau) + 0.000719 * math.cos(2*tau) + 0.000077 * math.sin(2 * tau)
-        solpars[j, 3] = 0.006918 - 0.399912 * math.cos(tau) + 0.070257 * math.sin(tau) - 0.006758 * math.cos(2*tau) + 0.000907 * math.sin(2 * tau) - 0.002697 * math.cos(3 * tau) + 0.00148 * math.sin(3 * tau)
-        solpars[j, 4] = solpars[j, 3] * 180 / math.pi
-        solpars[j, 5] = math.acos(-math.tan(lat * math.pi / 180) * math.tan(solpars[j, 3])) * 180 / math.pi
-        solpars[j, 6] = 2 / 15 * solpars[j, 5]
+        solparam[j, 0] = tau
+        solparam[j, 1] = tau * 180 / math.pi
+        solparam[j, 2] = 1.00011 + 0.034221 * math.cos(tau) + 0.00128 * math.sin(tau) + 0.000719 * math.cos(2*tau) + 0.000077 * math.sin(2 * tau)
+        solparam[j, 3] = 0.006918 - 0.399912 * math.cos(tau) + 0.070257 * math.sin(tau) - 0.006758 * math.cos(2*tau) + 0.000907 * math.sin(2 * tau) - 0.002697 * math.cos(3 * tau) + 0.00148 * math.sin(3 * tau)
+        solparam[j, 4] = solparam[j, 3] * 180 / math.pi
+        solparam[j, 5] = math.acos(-math.tan(lat * math.pi / 180) * math.tan(solparam[j, 3])) * 180 / math.pi
+        solparam[j, 6] = 2 / 15 * solparam[j, 5]
 
     # Duplicate line 365 for years with 366 days
-    solpars[365, :] = solpars[364, :]
+    solparam[365, :] = solparam[364, :]
     timecorr[365, :] = timecorr[364, :]
 
-    return solpars, timecorr
+    return solparam, timecorr
 
 
 def haversine(lat1, lon1, lat2, lon2):
@@ -262,7 +264,7 @@ def relshad(dem, mask, lats, lons, solh, sdirfn):
                 x_list = np.round(np.linspace(idx[0], idx[1], len(lon_list_short)))
 
                 # Calculate ALTITUDE along profile
-                zi = z[y_list.astype(np.int), x_list.astype(np.int)]
+                zi = z[y_list.astype(int), x_list.astype(int)]
 
                 # Calclulate DISTANCE along profile
                 d_list = []
@@ -292,7 +294,7 @@ def LUTshad(solpars, timecorr, lat, elvgrid, maskgrid, lats, lons, STEP, TCART):
            timecorr:  Time correction due to orbital forcing
            lat:       Latitude at AWS
            elvgrid:   DEM
-           maksgrid:  Glacier mask
+           maskgrid:  Glacier mask
            lats:      Latitudes
            lons:      Longitudes
            STEP:      Time step (s)
@@ -394,7 +396,7 @@ def calcRad(solPars, timecorr, doy, hour, lat, tempgrid, pgrid, rhgrid, cldgrid,
            rhgrid:    Relative Humidity
            cldgrid:   Cloud fraction
            elvgrid:   DEM
-           maksgrid:  Glacier mask
+           maskgrid:  Glacier mask
            slopegrid: Slope
            aspectgrid:Aspect
            shad1yr:   LUT topographic shading
@@ -447,8 +449,10 @@ def calcRad(solPars, timecorr, doy, hour, lat, tempgrid, pgrid, rhgrid, cldgrid,
         k_aes[k_aes > 1.0] = 1.0  # Aerosol factor: cannot be > 1
         TAUa = k_aes ** (mopt)
         TAUaa = 1.0 - (1.0 - alphss) * (1 - pgrid / 1013.25 * mopt + (pgrid / 1013.25 * mopt) ** (1.06)) * (1.0 - TAUa)
-        TAUw = 1.0 - 2.4959 * mopt * (46.5 * vp_interp / tempgrid) / ((1.0 + 79.034 * mopt * (46.5 * vp_interp / tempgrid)) **
-                                                                      0.6828 + 6.385 * mopt * (46.5 * vp_interp / tempgrid))
+        TAUw = 1.0 - 2.4959 * mopt * (46.5 * vp_interp / tempgrid) / (
+            (1.0 + 79.034 * mopt * (46.5 * vp_interp / tempgrid)) ** 0.6828
+            + 6.385 * mopt * (46.5 * vp_interp / tempgrid)
+        )
         taucs = TAUr * TAUg * TAUa * TAUw
 
         sdir = Sol0 * eccorr * sin_h * taucs  # Direct solar radiation on horizontal surface, clear-sky
@@ -457,27 +461,54 @@ def calcRad(solPars, timecorr, doy, hour, lat, tempgrid, pgrid, rhgrid, cldgrid,
         grcs = sdir + Dcs  # Potential clear-sky global radiation
 
         # Correction for slope and aspect (Iqbal 1983)
-        cos_zetap1 = (np.cos(np.radians(slopegrid)) * np.sin(np.radians(lat)) - np.cos(np.radians(lat)) * np.cos(np.radians(180 - aspectgrid)) *
-                      np.sin(np.radians(slopegrid))) * np.sin(soldec)
-        cos_zetap2 = (np.sin(np.radians(lat)) * np.cos(np.radians(180 - aspectgrid)) * np.sin(np.radians(slopegrid)) +
-                      np.cos(np.radians(slopegrid)) * np.cos(math.radians(lat))) * np.cos(soldec) * np.cos(stime * np.pi / 180)
-        cos_zetap3 = np.sin(np.radians(180 - aspectgrid)) * np.sin(np.radians(slopegrid)) * np.cos(soldec) * np.sin(stime * np.pi / 180)
+        slopegrid_rad = np.radians(slopegrid)  # avoid recalculating
+        cos_slopegrid = np.cos(slopegrid_rad)
+        sin_slopegrid = np.sin(slopegrid_rad)
+        rot_aspectgrid = np.radians(180 - aspectgrid)
+        lat_rad = np.radians(lat)
+        cos_zetap1 = (
+            cos_slopegrid * np.sin(lat_rad)
+            - np.cos(lat_rad) * np.cos(rot_aspectgrid) * sin_slopegrid
+        ) * np.sin(soldec)
+        cos_zetap2 = (
+            (
+                np.sin(lat_rad)
+                * np.cos(rot_aspectgrid)
+                * sin_slopegrid
+                + cos_slopegrid * np.cos(math.radians(lat))
+            )
+            * np.cos(soldec)
+            * np.cos(stime * np.pi / 180)
+        )
+        cos_zetap3 = (
+            np.sin(rot_aspectgrid)
+            * sin_slopegrid
+            * np.cos(soldec)
+            * np.sin(stime * np.pi / 180)
+        )
         cos_zetap = cos_zetap1 + cos_zetap2 + cos_zetap3
 
         # Clear-sky direct solar radiation at surface (aspect & slope corrected)
         swidir0 = Sol0 * eccorr * cos_zetap * taucs
         swidir0[cos_zetap < 0.0] = 0.0  # self-shaded cells set to 0
-        illu = elvgrid * 0.0
-        illu = shad1yr[np.int(((doy - 1) * (86400 / STEP)) + (hour / (STEP / 3600))), :, :]
+        # illu = elvgrid * 0.0
+        illu = shad1yr[int(((doy - 1) * (86400 / STEP)) + (hour / (STEP / 3600))), :, :]
         swidir0[illu == 0.0] = 0.0
         sdir[illu == 0.0] = 0.0
 
         # Correction for cloud fraction
-        swidiff[cldgrid > 0.0] = grcs[cldgrid > 0.0] * (((100 - Cf * 100) - dif1) / 100 * cldgrid[cldgrid > 0.0] +
-                                    (dif1 / 100)) * gridsvf[cldgrid > 0.0]  # diffuse amount as percentage of direct rad.
+        swidiff[cldgrid > 0.0] = (
+            grcs[cldgrid > 0.0]
+            * (
+                ((100 - Cf * 100) - dif1) / 100 * cldgrid[cldgrid > 0.0]
+                + (dif1 / 100)
+            )
+            * gridsvf[cldgrid > 0.0]
+        )  # diffuse amount as percentage of direct rad.
         swidiff[cldgrid == 0.0] = Dcs[cldgrid == 0.0] * gridsvf[cldgrid == 0.0]
-        swiasky[:, :] = swidir0 * (1 - (1 - dirovc) * cldgrid) + swidiff  # all-sky solar radiation at surface
-
+        # all-sky solar radiation at surface
+        swiasky[:, :] = swidir0 * (1 - (1 - dirovc) * cldgrid) + swidiff
+        
     else:
         TOAR = 0.0
         swiasky[maskgrid == 1] = 0 * elvgrid[maskgrid == 1]
diff --git a/cosipy/modules/roughness.py b/cosipy/modules/roughness.py
index 23595a22..0fab7ec4 100644
--- a/cosipy/modules/roughness.py
+++ b/cosipy/modules/roughness.py
@@ -8,7 +8,11 @@ def updateRoughness(GRID):
     if roughness_method == 'Moelg12':
         sigma = method_Moelg(GRID)
     else:
-        raise ValueError("Roghness method = \"{:s}\" is not allowed, must be one of {:s}".format(roughness_method,", ".join(roughness_allowed)))
+        error_message = (
+            f'Roughness method = "{roughness_method}" is not allowed,',
+            f'must be one of {", ".join(roughness_allowed)}'
+        )
+        raise ValueError(" ".join(error_message))
 
     return sigma
 
diff --git a/cosipy/modules/surfaceTemperature.py b/cosipy/modules/surfaceTemperature.py
index c7163314..9b0efa86 100644
--- a/cosipy/modules/surfaceTemperature.py
+++ b/cosipy/modules/surfaceTemperature.py
@@ -14,9 +14,9 @@ def update_surface_temperature(GRID, dt, z, z0, T2, rH2, p, SWnet, u2, RAIN, SLO
 
         GRID    ::  Grid structure
         T0      ::  Surface temperature [K]
-	    dt      ::  Integration time [s] -- can vary in WRF_X_CSPY
+        dt      ::  Integration time [s] -- can vary in WRF_X_CSPY
         alpha   ::  Albedo [-]
-	    z       ::  Measurement height [m] -- varies in WRF_X_CSPY
+        z       ::  Measurement height [m] -- varies in WRF_X_CSPY
         z0      ::  Roughness length [m]
         T2      ::  Air temperature [K]
         rH2     ::  Relative humidity [%]
@@ -39,7 +39,7 @@ def update_surface_temperature(GRID, dt, z, z0, T2, rH2, p, SWnet, u2, RAIN, SLO
         SWnet   ::  Shortwave radiation budget [W m^-2]
         rho     ::  Air density [kg m^-3]
         Lv      ::  Latent heat of vaporization [J kg^-1]
-	    MOL     ::  Monin-Obukhov length
+        MOL     ::  Monin-Obukhov length
         Cs_t    ::  Stanton number [-]
         Cs_q    ::  Dalton number [-]
         q0      ::  Mixing ratio at the surface [kg kg^-1]
@@ -58,7 +58,7 @@ def update_surface_temperature(GRID, dt, z, z0, T2, rH2, p, SWnet, u2, RAIN, SLO
         res = minimize(eb_optim, GRID.get_node_temperature(0), method=sfc_temperature_method,
                        bounds=((lower_bnd_ts, zero_temperature),),tol=1e-2,
                        args=(GRID, dt, z, z0, T2, rH2, p, SWnet, u2, RAIN, SLOPE, B_Ts, LWin, N))
-		       
+
     elif sfc_temperature_method == 'Newton':
         try:
             res = newton(eb_optim, np.array([GRID.get_node_temperature(0)]), tol=1e-2, maxiter=50,
@@ -66,10 +66,10 @@ def update_surface_temperature(GRID, dt, z, z0, T2, rH2, p, SWnet, u2, RAIN, SLO
             if res < lower_bnd_ts:
                 raise ValueError("TS Solution is out of bounds")
             res = SimpleNamespace(**{'x':min(np.array([zero_temperature]),res),'fun':None})
-	    
+
         except (RuntimeError,ValueError):
-             #Workaround for non-convergence and unboundedness
-             res = minimize(eb_optim, GRID.get_node_temperature(0), method='SLSQP',
+            #Workaround for non-convergence and unboundedness
+            res = minimize(eb_optim, GRID.get_node_temperature(0), method='SLSQP',
                        bounds=((lower_bnd_ts, zero_temperature),),tol=1e-2,
                        args=(GRID, dt, z, z0, T2, rH2, p, SWnet, u2, RAIN, SLOPE, B_Ts, LWin, N))
     else:
@@ -91,39 +91,48 @@ def update_surface_temperature(GRID, dt, z, z0, T2, rH2, p, SWnet, u2, RAIN, SLO
 
 
 @njit
-def interp_subT(GRID):
-    ''' Interpolate subsurface temperature to depths used for ground heat flux computation'''
-    
-    # Cumulative layer depths
-    layer_heights_cum = np.cumsum(np.array(GRID.get_height()))
+def get_subsurface_temperature(GRID, cumulative_depth: np.ndarray, zlt: float):
+    """Get subsurface temperature.
+
+    Args:
+        GRID (Grid): Gridded data instance.
+        cumulative_depth: Cumulative glacier layer heights [m].
+        zlt: Interpolation depth [m].
+
+    Returns:
+        float: Subsurface temperature at the interpolation depth.
+    """
 
     # Find indexes of two depths for temperature interpolation
-    idx1_depth_1 = np.abs(layer_heights_cum - zlt1).argmin()
-    depth_1 = layer_heights_cum.flat[np.abs(layer_heights_cum - zlt1).argmin()]
+    idx1_depth = np.abs(cumulative_depth - zlt).argmin()
+    depth = cumulative_depth.flat[idx1_depth]
 
-    if depth_1 > zlt1:
-        idx2_depth_1 = idx1_depth_1 - 1
+    if depth > zlt:
+        idx2_depth = idx1_depth - 1
     else:
-        idx2_depth_1 = idx1_depth_1 + 1
-    Tz1 = GRID.get_node_temperature(idx1_depth_1) + \
-		((GRID.get_node_temperature(idx1_depth_1) - GRID.get_node_temperature(idx2_depth_1)) / \
-            	(layer_heights_cum[idx1_depth_1] - layer_heights_cum[idx2_depth_1])) * \
-		(zlt1 - layer_heights_cum[idx1_depth_1])
+        idx2_depth = idx1_depth + 1
 
-    idx1_depth_2 = np.abs(layer_heights_cum - zlt2).argmin()
-    depth_2 = layer_heights_cum.flat[np.abs(layer_heights_cum - zlt2).argmin()]
+    temperature_idx1 = GRID.get_node_temperature(idx1_depth)
+    t_z = temperature_idx1 + (
+        (temperature_idx1 - GRID.get_node_temperature(idx2_depth))
+        / (cumulative_depth[idx1_depth] - cumulative_depth[idx2_depth])
+    ) * (zlt - cumulative_depth[idx1_depth])
+
+    return t_z
 
-    if depth_2 > zlt2:
-        idx2_depth_2 = idx1_depth_2 - 1
-    else:
-        idx2_depth_2 = idx1_depth_2 + 1
 
-    Tz2 = GRID.get_node_temperature(idx1_depth_2) + \
-		((GRID.get_node_temperature(idx1_depth_2) - GRID.get_node_temperature(idx2_depth_2)) / \
-        	(layer_heights_cum[idx1_depth_2] - layer_heights_cum[idx2_depth_2])) * \
-		(zlt2 - layer_heights_cum[idx1_depth_2])
-    return np.array([Tz1,Tz2])
+@njit
+def interp_subT(GRID) -> np.ndarray:
+    """Interpolate subsurface temperature to depths used for ground heat flux."""
     
+    # Cumulative layer depths
+    layer_heights_cum = np.cumsum(np.array(GRID.get_height()))
+
+    t_z1 = get_subsurface_temperature(GRID, layer_heights_cum, zlt1)
+    t_z2 = get_subsurface_temperature(GRID, layer_heights_cum, zlt2)
+
+    return np.array([t_z1, t_z2])
+
 
 @njit
 def eb_fluxes(GRID, T0, dt, z, z0, T2, rH2, p, u2, RAIN, SLOPE, B_Ts, LWin=None, N=None):
@@ -134,16 +143,15 @@ def eb_fluxes(GRID, T0, dt, z, z0, T2, rH2, p, u2, RAIN, SLOPE, B_Ts, LWin=None,
         GRID    ::  Grid structure
         T0      ::  Surface temperature [K]
         dt      ::  Integration time [s]
-        alpha   ::  Albedo [-]
-	z       ::  Measurement height [m]
+        z       ::  Measurement height [m]
         z0      ::  Roughness length [m]
         T2      ::  Air temperature [K]
         rH2     ::  Relative humidity [%]
         p       ::  Air pressure [hPa]
-        G       ::  Incoming shortwave radiation [W m^-2]
         u2      ::  Wind velocity [m S^-1]
         RAIN    ::  RAIN (mm)
         SLOPE   ::  Slope of the surface [degree]
+        B_Ts    ::  Subsurface temperatures at interpolation depths [K]
         LWin    ::  Incoming longwave radiation [W m^-2]
         N       ::  Fractional cloud cover [-]
 
@@ -158,7 +166,7 @@ def eb_fluxes(GRID, T0, dt, z, z0, T2, rH2, p, u2, RAIN, SLOPE, B_Ts, LWin=None,
         SWnet   ::  Shortwave radiation budget [W m^-2]
         rho     ::  Air density [kg m^-3]
         Lv      ::  Latent heat of vaporization [J kg^-1]
-	MOL     ::  Monin Obhukov length
+        MOL     ::  Monin Obhukov length
         Cs_t    ::  Stanton number [-]
         Cs_q    ::  Dalton number [-]
         q0      ::  Mixing ratio at the surface [kg kg^-1]
@@ -209,6 +217,10 @@ def eb_fluxes(GRID, T0, dt, z, z0, T2, rH2, p, u2, RAIN, SLOPE, B_Ts, LWin=None,
     z0q = z0/10     # Roughness length for moisture
     L = None
  
+    # Avoid recalculating
+    slope_radians = np.radians(SLOPE)
+    cos_slope_radians = np.cos(slope_radians)
+
     # Monin-Obukhov stability correction
     if stability_correction == 'MO':
         L = 0.0
@@ -223,19 +235,20 @@ def eb_fluxes(GRID, T0, dt, z, z0, T2, rH2, p, u2, RAIN, SLOPE, B_Ts, LWin=None,
             ust = ustar(u2,z,z0,L)
         
             # Sensible heat flux for neutral conditions
+            delta_phi_tq = phi_tq(z, L) - phi_tq(z0, L)
             Cd = np.power(0.41,2.0) / np.power(np.log(z/z0) - phi_m(z,L) - phi_m(z0,L),2.0)
-            Cs_t = 0.41*np.sqrt(Cd) / (np.log(z/z0t) - phi_tq(z,L) - phi_tq(z0,L))
-            Cs_q = 0.41*np.sqrt(Cd) / (np.log(z/z0q) - phi_tq(z,L) - phi_tq(z0,L))
+            Cs_t = 0.41*np.sqrt(Cd) / (np.log(z/z0t) - delta_phi_tq)
+            Cs_q = 0.41*np.sqrt(Cd) / (np.log(z/z0q) - delta_phi_tq)
         
             # Surface heat flux
-            H = rho * spec_heat_air * Cs_t * u2 * (T2-T0) * np.cos(np.radians(SLOPE))
+            H = rho * spec_heat_air * Cs_t * u2 * (T2-T0) * cos_slope_radians
         
             # Latent heat flux
-            LE = rho * Lv * Cs_q * u2 * (q2-q0) *  np.cos(np.radians(SLOPE))
+            LE = rho * Lv * Cs_q * u2 * (q2-q0) * cos_slope_radians
         
             # Monin-Obukhov length
             L = MO(rho, ust, T2, H)
-	    
+
             # Heat flux differences between iterations
             diff = np.abs(H0-H)
            
@@ -260,19 +273,20 @@ def eb_fluxes(GRID, T0, dt, z, z0, T2, rH2, p, u2, RAIN, SLOPE, B_Ts, LWin=None,
 
         # Stability correction
         phi = 1
-        if (Ri > 0.01) & (Ri <= 0.2):
+        if 0.01 < Ri <= 0.2:
             phi = np.power(1-5*Ri,2)
         elif Ri > 0.2:
             phi = 0
 
         # Sensible heat flux
-        H = rho * spec_heat_air * Cs_t * u2 * (T2-T0) * phi * np.cos(np.radians(SLOPE))
+        H = rho * spec_heat_air * Cs_t * u2 * (T2-T0) * phi * cos_slope_radians
         
         # Latent heat flux
-        LE = rho * Lv * Cs_q * u2 * (q2-q0) * phi * np.cos(np.radians(SLOPE))
+        LE = rho * Lv * Cs_q * u2 * (q2-q0) * phi * cos_slope_radians
 
     else:
-        raise ValueError("Stability correction",stability_correction,"is not supported")	#numba refuses to print str(list)
+        msg = f"Stability correction {stability_correction} is not supported."
+        raise ValueError(msg)
     
     # Outgoing longwave radiation
     Lo = -surface_emission_coeff * sigma * np.power(T0, 4.0)
@@ -295,33 +309,58 @@ def eb_fluxes(GRID, T0, dt, z, z0, T2, rH2, p, u2, RAIN, SLOPE, B_Ts, LWin=None,
 
 
 @njit
-def phi_m(z,L):
-    """ Stability function for the momentum flux.
+def phi_m_stable(z: float, L: float) -> float:
+    """Get integrated stability function for stable conditions.
+
+    Args:
+        z: Height, [m].
+        L: Obukhov length, [m].
+
+    Returns:
+        Stability function for momentum under stable conditions.
     """
-    if (L>0):
-        if ((z/L)>0.0) & ((z/L)<=1.0):
-            return (-5*z/L)
-        elif ((z/L)>1.0):
-            return (1-5) * (1+np.log(z/L)) - (z/L) 
-    elif L<0:
-        x = np.power((1-16*z/L),0.25)
-        return 2*np.log((1+x)/2.0) + np.log((1+np.power(x,2.0))/2.0) - 2*np.arctan(x) + np.pi/2.0
+    zeta = z / L
+    if (zeta > 0.0) & (zeta <= 1.0):  # weak stability
+        return -5 * zeta
+    elif zeta > 1.0:  # strong stability
+        return (1 - 5) * (1 + np.log(zeta)) - zeta
     else:
         return 0.0
 
 
 @njit
-def phi_tq(z,L):
-    """ Stability function for the heat and moisture flux.
+def phi_m(z: float, L: float) -> float:
+    """Get integrated stability function for the momentum flux.
+
+    Args:
+        z: Height, [m].
+        L: Obukhov length, [m].
+
+    Returns:
+        Integrated stability function for the momentum flux.
     """
-    if (L>0):
-        if ((z/L)>0.0) & ((z/L)<=1.0):
-            return (-5*z/L)
-        elif ((z/L)>1.0):
-            return (1-5) * (1+np.log(z/L)) - (z/L) 
-    elif L<0:
-        x = np.power((1-19.3*z/L),0.25)
-        return 2*np.log((1+np.power(x,2.0))/2.0)
+    if L > 0:
+        return phi_m_stable(z, L)
+    elif L < 0:
+        x = np.power((1 - 16 * z / L), 0.25)
+        return (
+            2 * np.log((1 + x) / 2.0)
+            + np.log((1 + np.power(x, 2.0)) / 2.0)
+            - 2 * np.arctan(x)
+            + np.pi / 2.0
+        )
+    else:
+        return 0.0
+
+
+@njit
+def phi_tq(z: float, L: float) -> float:
+    """Stability function for the heat and moisture flux."""
+    if L > 0:
+        return phi_m_stable(z, L)
+    elif L < 0:
+        x = np.power((1 - 19.3 * z / L), 0.25)
+        return 2 * np.log((1 + np.power(x, 2.0)) / 2.0)
     else:
         return 0.0
 
@@ -371,4 +410,3 @@ def method_EW_Sonntag(T):
         # over ice
         Ew = 6.112 * np.exp((22.46*(T-273.16)) / ((T-0.55)))
     return Ew
-
diff --git a/dev_requirements.txt b/dev_requirements.txt
new file mode 100644
index 00000000..0a468145
--- /dev/null
+++ b/dev_requirements.txt
@@ -0,0 +1,20 @@
+numpy>=1.17.3
+pandas>2.0
+scipy>=1.1.0
+xarray>=0.18.0
+netcdf4>=1.5.1.2
+distributed>=2022.3,<=2023.2
+dask
+dask-jobqueue
+numba
+metpy
+gdal
+nco
+cdo
+matplotlib
+cartopy
+vtk
+richdem
+coveralls
+codecov
+pytest
diff --git a/postprocessing/field_plots/README b/postprocessing/field_plots/README
index 9b6be9aa..a68b8752 100644
--- a/postprocessing/field_plots/README
+++ b/postprocessing/field_plots/README
@@ -1,4 +1,30 @@
+Create simple distributed plots of one timestep or of a daily mean.
 
-Create simple distributed plots of one timestep
+python plot_cosipy_fields.py -f nc_file -d timestamp
+python plot_cosipy_fields.py -f nc_file -d timestamp [-m] [-s] [-v variable] [-t <int>]
 
-python plot_cosipy_fields.py -f cosipy_outputfile -d timedate -t option
+Example:
+
+Create a contour plot of the daily mean for surface mass balance for the
+Zhadang test data, and save the plot to disk.
+
+python plot_cosipy_fields.py \
+    -f "../../data/output/Zhadang_ERA5_20090101-20090110.nc"
+    -d "2009-01-05" \
+    -v "surfMB" \
+    -t 1 \
+    -m \
+    -s
+
+Required arguments:
+    -f, --file <path>       Path to .nc file.
+    -d, --date <str>        Target date or timestamp.
+
+Optional switches:
+    -s, --save              Save plot. File name is automatically generated.
+                                Default False.
+    -m, --mean              Plot daily mean instead of timestep. Default False.
+
+Optional arguments:
+    -v, --var <str>         Variable to plot. If not set, creates a stamp plot.
+    -t, --type <int>        Set plot type. 1: contour, 2: mesh. Default 1.
\ No newline at end of file
diff --git a/postprocessing/field_plots/plot_cosipy_fields.py b/postprocessing/field_plots/plot_cosipy_fields.py
index 252fe108..4fe202bc 100644
--- a/postprocessing/field_plots/plot_cosipy_fields.py
+++ b/postprocessing/field_plots/plot_cosipy_fields.py
@@ -1,222 +1,501 @@
+"""
+Usage: ``src/plot_cosipy_fields.py [-h] [-f <nc_file>] [-d] [...] [-v]``
+"""
+
+import argparse
+import re
+
+import cartopy.crs as ccrs
 import matplotlib
-matplotlib.use("TkAgg")
 import matplotlib.pyplot as plt
-import os
 import numpy as np
 import xarray as xr
-import cartopy.crs as ccrs
-import argparse
 
-def plotMesh(filename, pdate, var=None):
-    """ This creates a simple plot showing the 2D fields"""
-    DATA = xr.open_dataset(filename)
-    
-    plt.figure(figsize=(20, 12))
-    print(DATA)
-
-    if var is None:
-        ax = plt.subplot(4,3,1,projection=ccrs.PlateCarree())
-        DATA.H.sel(time=pdate).plot.pcolormesh('lon', 'lat', ax=ax);
-        gl = ax.gridlines(crs=ccrs.PlateCarree(), draw_labels=True,
-                          linewidth=1, color='gray', alpha=0.1, linestyle='--')
-        gl.xlabels_top = None
-        gl.ylabels_right = False
-        
-        ax = plt.subplot(4,3,2,projection=ccrs.PlateCarree())
-        DATA.LE.sel(time=pdate).plot.pcolormesh('lon', 'lat', ax=ax);
-        gl = ax.gridlines(crs=ccrs.PlateCarree(), draw_labels=True,
-                          linewidth=1, color='gray', alpha=0.1, linestyle='--')
-        gl.xlabels_top = None
-        gl.ylabels_right = False
-        
-        ax = plt.subplot(4,3,3,projection=ccrs.PlateCarree())
-        DATA.B.sel(time=pdate).plot.pcolormesh('lon', 'lat', ax=ax);
-        gl = ax.gridlines(crs=ccrs.PlateCarree(), draw_labels=True,
-                          linewidth=1, color='gray', alpha=0.1, linestyle='--')
-        gl.xlabels_top = None
-        gl.ylabels_right = False
-        
-        ax = plt.subplot(4,3,4,projection=ccrs.PlateCarree())
-        DATA.SNOWHEIGHT.sel(time=pdate).plot.pcolormesh('lon', 'lat', ax=ax);
-        gl = ax.gridlines(crs=ccrs.PlateCarree(), draw_labels=True,
-                          linewidth=1, color='gray', alpha=0.1, linestyle='--')
-        gl.xlabels_top = None
-        gl.ylabels_right = False
-        
-        ax = plt.subplot(4,3,5,projection=ccrs.PlateCarree())
-        DATA.surfM.sel(time=pdate).plot.pcolormesh('lon', 'lat', ax=ax);
-        gl = ax.gridlines(crs=ccrs.PlateCarree(), draw_labels=True,
-                          linewidth=1, color='gray', alpha=0.1, linestyle='--')
-        gl.xlabels_top = None
-        gl.ylabels_right = False
-
-        ax = plt.subplot(4,3,6,projection=ccrs.PlateCarree())
-        DATA.TS.sel(time=pdate).plot.pcolormesh('lon', 'lat', ax=ax);
-        gl = ax.gridlines(crs=ccrs.PlateCarree(), draw_labels=True,
-                          linewidth=1, color='gray', alpha=0.1, linestyle='--')
-        gl.xlabels_top = None
-        gl.ylabels_right = False
-        
-        ax = plt.subplot(4,3,7,projection=ccrs.PlateCarree())
-        DATA.G.sel(time=pdate).plot.pcolormesh('lon', 'lat', ax=ax);
-        gl = ax.gridlines(crs=ccrs.PlateCarree(), draw_labels=True,
-                          linewidth=1, color='gray', alpha=0.1, linestyle='--')
-        gl.xlabels_top = None
-        gl.ylabels_right = False
-        
-        ax = plt.subplot(4,3,8,projection=ccrs.PlateCarree())
-        DATA.Q.sel(time=pdate).plot.pcolormesh('lon', 'lat', ax=ax);
-        gl = ax.gridlines(crs=ccrs.PlateCarree(), draw_labels=True,
-                          linewidth=1, color='gray', alpha=0.1, linestyle='--')
-        gl.xlabels_top = None
-        gl.ylabels_right = False
-        
-        ax = plt.subplot(4,3,9,projection=ccrs.PlateCarree())
-        DATA.LWout.sel(time=pdate).plot.pcolormesh('lon', 'lat', ax=ax);
-        gl = ax.gridlines(crs=ccrs.PlateCarree(), draw_labels=True,
-                          linewidth=1, color='gray', alpha=0.1, linestyle='--')
-        gl.xlabels_top = None
-        gl.ylabels_right = False
-        
-        ax = plt.subplot(4,3,10,projection=ccrs.PlateCarree())
-        DATA.surfMB.sel(time=pdate).plot.pcolormesh('lon', 'lat', ax=ax);
-        gl = ax.gridlines(crs=ccrs.PlateCarree(), draw_labels=True,
-                          linewidth=1, color='gray', alpha=0.1, linestyle='--')
-        gl.xlabels_top = None
-        gl.ylabels_right = False
-        
-        ax = plt.subplot(4,3,11,projection=ccrs.PlateCarree())
-        DATA.MB.sel(time=pdate).plot.pcolormesh('lon', 'lat', ax=ax);
-        gl = ax.gridlines(crs=ccrs.PlateCarree(), draw_labels=True,
-                          linewidth=1, color='gray', alpha=0.1, linestyle='--')
-        gl.xlabels_top = None
-        gl.ylabels_right = False
-        
-        ax = plt.subplot(4,3,12,projection=ccrs.PlateCarree())
-        DATA.REFREEZE.sel(time=pdate).plot.pcolormesh('lon', 'lat', ax=ax);
-        gl = ax.gridlines(crs=ccrs.PlateCarree(), draw_labels=True,
-                          linewidth=1, color='gray', alpha=0.1, linestyle='--')
-        gl.xlabels_top = None
-        gl.ylabels_right = False
-
-    
-    ax.coastlines()
-    plt.show()
-
-def plotContour(filename, pdate, var=None):
-    """ This creates a simple plot showing the 2D fields"""
-    DATA = xr.open_dataset(filename)
-    
-    plt.figure(figsize=(20, 12))
-    print(DATA)
-
-    if var is None:
-        ax = plt.subplot(4,3,1,projection=ccrs.PlateCarree())
-        DATA.H.sel(time=pdate).plot.contourf('lon', 'lat', ax=ax);
-        gl = ax.gridlines(crs=ccrs.PlateCarree(), draw_labels=True,
-                          linewidth=1, color='gray', alpha=0.1, linestyle='--')
-        gl.xlabels_top = None
-        gl.ylabels_right = False
-        
-        ax = plt.subplot(4,3,2,projection=ccrs.PlateCarree())
-        DATA.LE.sel(time=pdate).plot.contourf('lon', 'lat', ax=ax);
-        gl = ax.gridlines(crs=ccrs.PlateCarree(), draw_labels=True,
-                          linewidth=1, color='gray', alpha=0.1, linestyle='--')
-        gl.xlabels_top = None
-        gl.ylabels_right = False
-        
-        ax = plt.subplot(4,3,3,projection=ccrs.PlateCarree())
-        DATA.B.sel(time=pdate).plot.contourf('lon', 'lat', ax=ax);
-        gl = ax.gridlines(crs=ccrs.PlateCarree(), draw_labels=True,
-                          linewidth=1, color='gray', alpha=0.1, linestyle='--')
-        gl.xlabels_top = None
-        gl.ylabels_right = False
-        
-        ax = plt.subplot(4,3,4,projection=ccrs.PlateCarree())
-        DATA.SNOWHEIGHT.sel(time=pdate).plot.contourf('lon', 'lat', ax=ax);
-        gl = ax.gridlines(crs=ccrs.PlateCarree(), draw_labels=True,
-                          linewidth=1, color='gray', alpha=0.1, linestyle='--')
-        gl.xlabels_top = None
-        gl.ylabels_right = False
-        
-        ax = plt.subplot(4,3,5,projection=ccrs.PlateCarree())
-        DATA.surfM.sel(time=pdate).plot.contourf('lon', 'lat', ax=ax);
-        gl = ax.gridlines(crs=ccrs.PlateCarree(), draw_labels=True,
-                          linewidth=1, color='gray', alpha=0.1, linestyle='--')
-        gl.xlabels_top = None
-        gl.ylabels_right = False
-
-        ax = plt.subplot(4,3,6,projection=ccrs.PlateCarree())
-        DATA.TS.sel(time=pdate).plot.contourf('lon', 'lat', ax=ax);
-        gl = ax.gridlines(crs=ccrs.PlateCarree(), draw_labels=True,
-                          linewidth=1, color='gray', alpha=0.1, linestyle='--')
-        gl.xlabels_top = None
-        gl.ylabels_right = False
-        
-        ax = plt.subplot(4,3,7,projection=ccrs.PlateCarree())
-        DATA.G.sel(time=pdate).plot.contourf('lon', 'lat', ax=ax);
-        gl = ax.gridlines(crs=ccrs.PlateCarree(), draw_labels=True,
-                          linewidth=1, color='gray', alpha=0.1, linestyle='--')
-        gl.xlabels_top = None
-        gl.ylabels_right = False
-        
-        ax = plt.subplot(4,3,8,projection=ccrs.PlateCarree())
-        DATA.LWin.sel(time=pdate).plot.contourf('lon', 'lat', ax=ax);
-        gl = ax.gridlines(crs=ccrs.PlateCarree(), draw_labels=True,
-                          linewidth=1, color='gray', alpha=0.1, linestyle='--')
-        gl.xlabels_top = None
-        gl.ylabels_right = False
-        
-        ax = plt.subplot(4,3,9,projection=ccrs.PlateCarree())
-        DATA.LWout.sel(time=pdate).plot.contourf('lon', 'lat', ax=ax);
-        gl = ax.gridlines(crs=ccrs.PlateCarree(), draw_labels=True,
-                          linewidth=1, color='gray', alpha=0.1, linestyle='--')
-        gl.xlabels_top = None
-        gl.ylabels_right = False
-        
-        ax = plt.subplot(4,3,10,projection=ccrs.PlateCarree())
-        DATA.MB.sel(time=pdate).plot.contourf('lon', 'lat', ax=ax);
-        gl = ax.gridlines(crs=ccrs.PlateCarree(), draw_labels=True,
-                          linewidth=1, color='gray', alpha=0.1, linestyle='--')
-        gl.xlabels_top = None
-        gl.ylabels_right = False
-        
-        ax = plt.subplot(4,3,11,projection=ccrs.PlateCarree())
-        DATA.surfMB.sel(time=pdate).plot.contourf('lon', 'lat', ax=ax);
-        gl = ax.gridlines(crs=ccrs.PlateCarree(), draw_labels=True,
-                          linewidth=1, color='gray', alpha=0.1, linestyle='--')
-        gl.xlabels_top = None
-        gl.ylabels_right = False
-        
-        ax = plt.subplot(4,3,12,projection=ccrs.PlateCarree())
-        DATA.Q.sel(time=pdate).plot.contourf('lon', 'lat', ax=ax);
-        gl = ax.gridlines(crs=ccrs.PlateCarree(), draw_labels=True,
-                          linewidth=1, color='gray', alpha=0.1, linestyle='--')
-        gl.xlabels_top = None
-        gl.ylabels_right = False
-
-
-        #DATA[var].sel(time=pdate).plot.contourf('lon', 'lat', ax=ax);
-        
-    
-    ax.coastlines()
-    plt.show()
+matplotlib.use("TkAgg")
 
 
-if __name__ == "__main__":
-    
-    parser = argparse.ArgumentParser(description='Quick plot of the results file.')
-    parser.add_argument('-f', '-file', dest='file', help='File name')
-    parser.add_argument('-d', '-date', dest='pdate', help='Which date to plot')
-    parser.add_argument('-v', '-var', dest='var', help='Which variable')
-    parser.add_argument('-t', '-type', dest='type', help='(1) contour, (2) mesh', type=int)
-
-    args = parser.parse_args()
-    print(args.type)
-    if args.type==1:
-        print('Contour')
-        plotContour(args.file, args.pdate, args.var) 
+def check_2d(array: xr.DataArray):
+    """Checks if an input array has 2D spatial coordinates.
+
+    Raises:
+        ValueError: Spatial coordinates are not 2D.
+    """
+
+    for dimension in array.dims:
+        if dimension not in ["time", "layer"] and (array.dims[dimension]) <= 1:
+            raise ValueError("Spatial coordinates are not 2D.")
+
+
+def get_selection(
+    array: xr.DataArray, timestamp: str, mean: bool = False
+) -> xr.DataArray:
+    """Selects data from array at specific time or as a daily mean.
+
+    Args:
+        array: Labelled data array.
+        timestamp: Date or datetime index of target data.
+        mean: If True, computes and selects the daily mean. Otherwise,
+            selects data at ``timestamp``. Default False.
+
+    Returns:
+        Array selection at target time.
+
+    Raises:
+        KeyError: Selected time index refers to an entire day, not a
+            single timestep. Use ``--mean`` to plot the daily mean.
+    """
+
+    if not mean:
+        data = array.sel(time=timestamp)
+        if data.time.size > 1:
+            error_message = (
+                "Selected time index refers to an entire day,",
+                "not a single timestep.",
+                "Use `--mean` to plot the daily mean.",
+            )
+            raise KeyError(" ".join(error_message))
     else:
-        print('Mesh')
-        plotMesh(args.file, args.pdate, args.var) 
-    
+        data = (
+            array.resample(time="1D", skipna=True)
+            .mean()
+            .sel(time=timestamp, method="nearest")
+        )
+
+    return data
+
+
+def set_contour_format(x: float) -> str:
+    """Formats float into contour label."""
+
+    if round(x, 1) % 1:
+        elevation = f"{x:.1f}"
+    else:
+        elevation = f"{int(x)}"
+
+    if plt.rcParams["text.usetex"]:
+        label = rf"{elevation} m"
+    else:
+        label = f"{elevation} m"
+
+    return label
+
+
+def set_gridlines(ax: plt.Axes) -> plt.Axes:
+    """Projects to PlateCarrée, adds gridlines, and formats labels."""
+
+    gridlines = ax.gridlines(
+        crs=ccrs.PlateCarree(),
+        draw_labels=True,
+        linewidth=0.5,
+        color="gray",
+        alpha=0.1,
+        linestyle="--",
+    )
+    gridlines.top_labels = None
+    gridlines.right_labels = False
+
+    return ax
+
+
+def save_image(
+    figure: plt.Figure,
+    timestamp: str,
+    name: str = "",
+    suffix: str = "",
+    img_format="png",
+):
+    """Automatically generates file name, and saves image to disk.
+
+    Args:
+        figure: Figure object that is saved to disk.
+        timestamp: Date or datetime index of target data.
+        name: Name of variable. Default empty string.
+        suffix: Additional string appended to file name e.g. the type of
+            plot. Default empty string.
+        img_format: Output image format. Default "png".
+    """
+
+    if not isinstance(timestamp, str):
+        img_id = timestamp.strftime("%Y%m%d")
+    else:
+        img_id = timestamp
+    img_id = re.sub(r"\W+", "", str(img_id))  # avoid illegal file names
+
+    if name is None:
+        img_id = f"{img_id}_stamp_plot"
+    elif name:
+        img_id = f"{img_id}_{name}"
+
+    if suffix:
+        suffix = f"_{suffix}"
+
+    figure.savefig(fname=f"{img_id}{suffix}.{img_format}")
+
+
+def plot_topography(ax: plt.Axes, elevation: xr.DataArray) -> plt.Axes:
+    """Plots elevation contours with labels onto axis."""
+
+    topography = ax.contour(
+        elevation.lon, elevation.lat, elevation, colors="gray", linewidths=0.5
+    )
+    ax.clabel(
+        topography,
+        topography.levels,
+        inline=True,
+        fmt=set_contour_format,
+        fontsize=10,
+        colors="black",
+    )
+
+    return ax
+
+
+def plot_axes(
+    ax: plt.Axes,
+    array: xr.DataArray,
+    timestamp: str,
+    topography: xr.DataArray = None,
+    plot_type: str = "contour",
+    mean: bool = False,
+) -> plt.Axes:
+    """Plots data for a specific time or as a daily mean.
+
+    Args:
+        ax: Target axis.
+        array: XYZ data for a single variable.
+        timestamp: Date or datetime index of target data.
+        topography: Optional elevation data to plot contours. Default
+            None.
+        plot_type: Plot data as "contour" or "mesh". Default "contour".
+        mean: Plot daily mean instead of specific time. Default False.
+
+    Returns:
+        Labelled contour or mesh plot of variable, with optional
+        contours for topography.
+
+    Raises:
+        ValueError: <plot_type> is an invalid plot type. Use 'mesh' or
+            'contour'.
+    """
+
+    name = array.attrs.get("long_name", None)
+    unit_label = array.attrs.get("units", None)
+    c_map = matplotlib.colormaps["viridis"]  # coolwarm
+
+    if name.lower() == "elevation":
+        data = array  # HGT has no time dimensions
+        c_map = matplotlib.colormaps["gist_earth"]
+    else:
+        data = get_selection(array=array, timestamp=timestamp, mean=mean)
+
+    if topography is not None:
+        plot_topography(ax=ax, elevation=topography)
+
+    if plot_type.lower() == "contour":
+        data.plot.contourf(
+            "lon", "lat", ax=ax, cbar_kwargs={"label": unit_label}, cmap=c_map
+        )
+    elif plot_type.lower() == "mesh":
+        data.plot.pcolormesh(
+            "lon", "lat", ax=ax, cbar_kwargs={"label": unit_label}, cmap=c_map
+        )
+    else:
+        raise ValueError(
+            f"{plot_type} is an invalid plot type. Use 'mesh' or 'contour'."
+        )
+    ax = set_gridlines(ax=ax)
+    ax.set_title(name.title())
+
+    return ax
+
+
+def create_subplots_for_variables(variables: list = None) -> tuple:
+    """Creates adaptive subplots for an array of variables.
+
+    Array is shaped as [nrows[ncols]].
+
+    Returns:
+        tuple[plt.Figure, np.ndarray]: Figure with subplots for each
+        variable.
+
+    Raises:
+        ValueError: Passed empty variable array to adaptive subplots.
+    """
+
+    if not variables:
+        raise ValueError("Passed empty variable array to adaptive subplots.")
+    else:
+        plot_array = np.array(variables)
+        figure, axes = plt.subplots(
+            plot_array.shape[0],
+            plot_array.shape[1],
+            sharex="col",
+            sharey="row",
+            figsize=(30, 18),
+            subplot_kw=dict(projection=ccrs.PlateCarree()),
+        )
+
+    return figure, axes
+
+
+def plot_data(
+    data: xr.DataArray,
+    timestamp: str,
+    short_name: str = None,
+    mean: bool = False,
+    variables: list = None,
+    plot_type: str = "contour",
+) -> tuple:
+    """Creates stamp plots or a single plot.
+
+    Args:
+        data: Labelled data array.
+        timestamp: Datetime index of target data.
+        short_name: Short name of variable to plot. If None, a stamp
+            plot of several variables is produced. Default None.
+        mean: Plot daily mean instead of specific time. Default False.
+        variables: Array of variable names of shape [nrows[ncols]].
+            Default None.
+        plot_type: Plot data as "contour" or "mesh". Default "contour".
+    """
+
+    if short_name is None:  # stamp plots
+        _, axes = create_subplots_for_variables(variables)
+        for i in range(len(variables)):
+            for j, key in enumerate(variables[i]):
+                axes[i][j] = plot_axes(
+                    ax=axes[i][j],
+                    array=data[key],
+                    timestamp=timestamp,
+                    mean=mean,
+                    topography=data.HGT,
+                    plot_type=plot_type,
+                )
+    else:
+        plt.figure(figsize=(18, 10))
+        axes = plt.subplot(1, 1, 1, projection=ccrs.PlateCarree())
+        axes = plot_axes(
+            ax=axes,
+            array=data[short_name],
+            timestamp=timestamp,
+            mean=mean,
+            topography=data.HGT,
+            plot_type=plot_type,
+        )
+    figure = plt.gcf()
+
+    return figure, axes
+
+
+def plotMesh(
+    filename: str,
+    pdate: str,
+    var: str = None,
+    mean: bool = False,
+    save: bool = False,
+):
+    """This creates a simple mesh plot of the 2D fields.
+
+    Args:
+        filename: Path to nc data array.
+        pdate: Date or datetime index of target data.
+        var: Short name of variable to plot. If None, a stamp plot of
+            several variables is produced. Default None.
+        mean: Plot daily mean instead of specific time. Default False.
+        save: Save plot to disk. Default False.
+
+    Raises:
+        ValueError: Input data is not 2D.
+    """
+
+    data = xr.open_dataset(filename)
+    check_2d(data)
+
+    print(data)
+    variables = [  # structure of stamp plot
+        ["H", "LE", "B"],
+        ["SNOWHEIGHT", "surfM", "TS"],
+        ["G", "Q", "LWout"],
+        ["surfMB", "MB", "REFREEZE"],
+    ]
+
+    figure, _ = plot_data(
+        data=data,
+        variables=variables,
+        timestamp=pdate,
+        short_name=var,
+        mean=mean,
+        plot_type="mesh",
+    )
+    if save:
+        save_image(figure=figure, timestamp=pdate, name=var, suffix="mesh")
+    else:
+        plt.show()
+
+
+def plotContour(
+    filename: str,
+    pdate: str,
+    var: str = None,
+    mean: bool = False,
+    save: bool = False,
+):
+    """This creates a simple contour plot of the 2D fields.
+
+    Args:
+        filename: Path to nc data array.
+        pdate: Date or datetime index of target data.
+        var: Short name of variable to plot. If None, a stamp plot of
+            several variables is produced. Default None.
+        mean: Plot daily mean instead of specific time. Default False.
+        save: Save plot to disk. Default False.
+
+    Raises:
+        ValueError: Input data is not 2D.
+    """
+
+    data = xr.open_dataset(filename)
+    check_2d(data)
+
+    print(data)
+    variables = [  # structure of stamp plot
+        ["H", "LE", "B"],
+        ["SNOWHEIGHT", "surfM", "TS"],
+        ["G", "LWin", "LWout"],
+        ["MB", "surfMB", "Q"],
+    ]
+
+    figure, _ = plot_data(
+        data=data,
+        variables=variables,
+        timestamp=pdate,
+        short_name=var,
+        mean=mean,
+        plot_type="contour",
+    )
+    if save:
+        save_image(figure=figure, timestamp=pdate, name=var, suffix="contour")
+    else:
+        plt.show()
+
+
+def parse_arguments() -> argparse.Namespace:
+    """Parse user arguments.
+
+    Required arguments:
+        -f, --file <path>       Path to .nc file.
+        -d, --date <str>        Target date or timestamp.
+
+    Optional switches:
+        -h, --help              Show this help message and exit.
+        -s, --save              Save plot. File name is automatically
+                                    generated. Default False.
+        -m, --mean              Plot daily mean instead of timestep.
+                                    Default False.
+
+    Optional arguments:
+        -v, --var <str>         Variable to plot. If not set, creates a
+                                    stamp plot.
+        -t, --type <int>        Set plot type. 1: contour, 2: mesh.
+                                    Default 1.
+    """
+
+    tagline = "Plot results for single or all variables."
+    parser = argparse.ArgumentParser(prog="plot_cosipy_fields.py", description=tagline)
+    # Required
+    parser.add_argument(
+        "-f",
+        "--file",
+        dest="file",
+        required=True,
+        default=None,
+        type=str,
+        metavar="<path>",
+        help="Path to .nc file",
+    )
+    parser.add_argument(
+        "-d",
+        "--date",
+        dest="pdate",
+        type=str,
+        required=True,
+        default=None,
+        metavar="<str>",
+        help="Target date or timestamp",
+    )
+    # switches
+    parser.add_argument(
+        "-m",
+        "--mean",
+        dest="mean",
+        action="store_true",
+        help="Plot daily mean instead of timestep. Default False",
+    )
+    parser.add_argument(
+        "-s",
+        "--save",
+        dest="save",
+        action="store_true",
+        help="Save plot. File name is automatically generated. Default False.",
+    )
+    # Optional
+    parser.add_argument(
+        "-v",
+        "--var",
+        dest="variable",
+        type=str,
+        default=None,
+        metavar="<str>",
+        help="Variable to plot. If not set, creates a stamp plot",
+    )
+    parser.add_argument(
+        "-t",
+        "--type",
+        dest="plot_type",
+        type=int,
+        default=1,
+        choices=[1, 2, None],
+        metavar="<int>",
+        help="Set plot type. 1: contour, 2: mesh. Default 1",
+    )
+    arguments = parser.parse_args()
+
+    return arguments
+
+
+def main():
+    """Produce and save plots.
+
+    Produces field plots of single timesteps or daily mean for a single
+    or all variables. Passing a date without ``--mean`` leads to an
+    error.
+
+    Required arguments:
+        -i, --input <path>      Path to .nc file.
+        -d, --date <str>        Target date or timestamp.
+
+    Optional switches:
+        -h, --help              Show this help message and exit.
+        -s, --save              Save plot. File name is automatically
+                                    generated. Default False.
+        -m, --mean              Plot daily mean instead of timestep.
+                                    Default False.
+
+    Optional arguments:
+        -v, --var <str>         Variable to plot. If not set, creates a
+                                    stamp plot.
+        -t, --type <int>        Set plot type. 1: contour, 2: mesh.
+                                    Default 1.
+    """
+
+    args = parse_arguments()
+    if args.plot_type == 1:
+        print("Contour")
+        plotContour(args.file, args.pdate, args.variable, args.mean, args.save)
+    else:
+        print("Mesh")
+        plotMesh(args.file, args.pdate, args.variable, args.mean, args.save)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/postprocessing/vtk_plots/README b/postprocessing/vtk_plots/README
new file mode 100644
index 00000000..dcdc89f3
--- /dev/null
+++ b/postprocessing/vtk_plots/README
@@ -0,0 +1,27 @@
+Create a 3D model of glacier surface for one timestep or for a daily mean.
+
+python plot_cosipy_fields_vtk.py -i nc_file -t timestamp -n variable
+python plot_cosipy_fields_vtk.py -i nc_file -t timestamp [-n name] [-m] 
+
+Example:
+
+Create a 3D model of the daily mean for surface mass balance for the
+Zhadang test data, and save the plot as an image.
+
+python plot_cosipy_fields_vtk.py \
+  -i "../../data/output/Zhadang_ERA5_20090101-20090110.nc" \
+  -t "2009-01-05" \
+  -n "surfMB" \
+  -m
+    
+Required arguments:
+    -i, --input <path>      Path to .nc file.
+    -t, --time <str>        Target date or timestamp.
+    -n, --name <str>        Name of variable to plot as contours.
+
+Optional switches:
+    -h, --help              Show this help message and exit.
+    -m, --mean              Plot daily mean instead of timestep.
+                                Default False.
+    -g, --gltf              Export plot as a 3D glTF object.
+                                Default False.
diff --git a/postprocessing/vtk_plots/plot_cosipy_fields_vtk.py b/postprocessing/vtk_plots/plot_cosipy_fields_vtk.py
index e385ba02..a9265fb8 100644
--- a/postprocessing/vtk_plots/plot_cosipy_fields_vtk.py
+++ b/postprocessing/vtk_plots/plot_cosipy_fields_vtk.py
@@ -1,232 +1,553 @@
-import vtk
+import argparse
+import re
+import warnings
+from itertools import product
+
 import numpy as np
-from vtk.util import numpy_support
-import math
-from scipy.optimize import minimize
-from scipy.optimize import curve_fit
-from scipy.optimize import fmin
-#import matplotlib.pyplot as plt
+import vtk
 import xarray as xr
-from itertools import product
+from vtk.util import numpy_support
+
 
 def main():
-    #createDEM_v1()
-    #createDEM_v2()
-    add_scalar('TS', '2017-07-01 12:00:00')
+    args = parse_arguments()
+
+    createDEM_v1(file_path=args.input_file)
+    # createDEM_v2(file_path=args.input_file)
+    add_scalar(
+        file_path=args.input_file,
+        var=args.name,
+        timestamp=args.timestamp,
+        mean=args.mean,
+        gltf=args.gltf,
+    )
+
+
+def check_2d(array: xr.DataArray):
+    """Checks if an input array has 2D spatial coordinates.
+
+    Raises:
+        ValueError: Spatial coordinates are not 2D.
+    """
+
+    for dimension in array.dims:
+        if dimension not in ["time", "layer"] and (array.dims[dimension]) <= 1:
+            raise ValueError("Spatial coordinates are not 2D.")
+
+
+def get_selection(
+    array: xr.Dataset, timestamp: str, mean: bool = False
+) -> xr.DataArray:
+    """Selects array from dataset at specific time or as a daily mean.
+
+    Args:
+        array (xr.Dataset or xr.DataArray): Labelled data.
+        timestamp: Time index of target data.
+        mean: If True, computes and selects the daily mean. Otherwise,
+            selects data at ``timestamp``. Default False.
+
+    Returns:
+        Array selection at target time.
+    """
+
+    if not mean:
+        data = array.sel(time=timestamp, method="nearest")
+    else:
+        data = (
+            array.resample(time="1D", skipna=True)
+            .mean()
+            .sel(time=timestamp, method="nearest")
+        )
+
+    return data
+
+
+def get_coords(data) -> tuple:
+    """Get latitude/longitude keys.
+
+    Args:
+        data (xr.Dataset or xr.DataArray): Labelled dataset.
 
-def createDEM_v1():
+    Returns:
+        tuple[str, str]: Keys referring to latitude and longitude.
+    """
 
-    ds = xr.open_dataset('../data/output/Peru_20160601-20180530_comp4.nc')
-    
+    if "south_north" in data.keys():
+        northing = "south_north"
+        easting = "west_east"
+    else:
+        northing = "lat"
+        easting = "lon"
+
+    return northing, easting
+
+
+def create_points(data: xr.Dataset) -> vtk.vtkPoints:
+    """Convert xarray dataset to points."""
+
+    print("Writing points...\n")
+    latitude, longitude = get_coords(data)
     points = vtk.vtkPoints()
-    
-    numPoints = ds.south_north.size*ds.west_east.size
-   
-    print('Write points \n')
-    for i,j in product(ds.south_north.values,ds.west_east.values):
-            points.InsertNextPoint(ds.lat.isel(south_north=i,west_east=j), ds.lon.isel(south_north=i,west_east=j), ds.HGT.isel(south_north=i,west_east=j).values/6370000.0)
-    
-    print('Create unstructured grid \n') 
+    if "south_north" in data.keys():
+        for i, j in product(data[longitude].values, data[latitude].values):
+            points.InsertNextPoint(
+                data[longitude].sel(west_east=i),
+                data[latitude].sel(south_north=j),
+                data.HGT.sel(west_east=i, south_north=j).values / 6370000.0,
+            )
+    else:
+        for i, j in product(data[longitude].values, data[latitude].values):
+            points.InsertNextPoint(
+                data[longitude].sel(lon=i),
+                data[latitude].sel(lat=j),
+                data.HGT.sel(lon=i, lat=j).values / 6370000.0,
+            )
+
+    return points
+
+
+def create_unstructured_grid(
+    points: vtk.vtkPoints, cells: vtk.vtkCellArray = None
+) -> vtk.vtkUnstructuredGrid:
+    """Creates unstructured grid from points and/or cells."""
+
+    print("Create unstructured grid...\n")
+    grid = vtk.vtkUnstructuredGrid()
+    if points is not None:
+        grid.SetPoints(points)
+    if cells is not None:
+        grid.SetCells(vtk.VTK_QUAD, cells)
+
+    return grid
+
+
+def write_unstructured_grid_to_file(
+    unstructured_grid: vtk.vtkUnstructuredGrid, file_name: str = "cosipy.vtu"
+):
+    """Writes unstructured grid data to .vtu file."""
+
+    writer = vtk.vtkXMLUnstructuredGridWriter()
+    writer.SetFileName(file_name)
+    writer.SetInputData(unstructured_grid)
+    writer.Write()
+
+
+def write_render_to_gltf(
+    render: vtk.vtkRenderer,
+    variable: str,
+    timestamp: str,
+):
+    """Writes rendered image to a 3D .gltf object.
+
+    Raises:
+        AttributeError: GLTF export is only supported in VTK >9.2.
+    """
+
+    print("Saving 3D object...\n")
+    output_file = set_filename(name=variable, timestamp=timestamp, fmt="gltf")
+    try:
+        writer = vtk.vtkGLTFExporter()
+        writer.InlineDataOn()
+        writer.SetActiveRenderer(render)
+        writer.SetRenderWindow(render.GetRenderWindow())
+        writer.SetFileName(output_file)
+        writer.Update()
+        writer.Write()
+        print("3D object saved.\n")
+    except AttributeError:
+        print("GLTF export is only supported in VTK >9.2.")
+
+
+def write_render_to_image(
+    render: vtk.vtkRenderLargeImage, variable: str, timestamp: str
+):
+    """Writes rendered image to .png file."""
+
+    print("Saving image...\n")
+    writer = vtk.vtkPNGWriter()
+    writer.SetInputConnection(render.GetOutputPort())
+    output_file = set_filename(name=variable, timestamp=timestamp)
+    writer.SetFileName(output_file)
+    writer.Write()
+    print("Image saved.\n")
+
+
+def smooth_mesh(mesh: vtk.vtkPolyData) -> vtk.vtkButterflySubdivisionFilter:
+    """Smooths mesh using butterfly subdivision."""
+
+    subdivision = vtk.vtkButterflySubdivisionFilter()
+    subdivision.SetInputData(mesh)
+    subdivision.Update()
+
+    # smoother = vtk.vtkWindowedSincPolyDataFilter()
+    # smoother.SetInputConnection(delaunay.GetOutputPort())
+    # smoother.SetNumberOfIterations(5)
+    # smoother.BoundarySmoothingOff()
+    # smoother.FeatureEdgeSmoothingOff()
+    # smoother.SetFeatureAngle(120.0)
+    # smoother.SetPassBand(0.001)
+    # smoother.NonManifoldSmoothingOff()
+    # smoother.NormalizeCoordinatesOff()
+    # smoother.Update()
+
+    return subdivision
+
+
+def createDEM_v1(
+    file_path: str = "../../data/output/Zhadang_ERA5_20090101-20090110.nc",
+):
+    """Map dataset to DEM using Delaunay triangulation.
+
+    The DEM is saved as a vtu file.
+
+    .. todo:: add terrain smoothing.
+
+    .. note:: smoothing the DEM on creation means add_scalar will not
+        find the closest points, as the planes no longer intersect. If
+        smoothing occurs at the start of plotSurface, the data lost
+        through triangulating is not recovered.
+    """
+
+    ds = xr.open_dataset(file_path)
+    check_2d(ds)
+    print("Creating DEM...\n")
+
+    points = create_points(data=ds)
     polydata = vtk.vtkPolyData()
     polydata.SetPoints(points)
 
-    delaunay = vtk.vtkDelaunay2D()
+    delaunay = vtk.vtkDelaunay3D()
     delaunay.SetInputData(polydata)
     delaunay.Update()
 
-#    subdivision = vtk.vtkButterflySubdivisionFilter()
-#    subdivision.SetInputConnection(delaunay.GetOutputPort())
-#    subdivision.Update()
+    append_filter = vtk.vtkAppendFilter()
+    append_filter.AddInputData(delaunay.GetOutput())
+    append_filter.Update()
 
-    #smoother = vtk.vtkWindowedSincPolyDataFilter()
-    #smoother.SetInputConnection(delaunay.GetOutputPort())
-    #smoother.SetNumberOfIterations(5)
-    #smoother.BoundarySmoothingOff()
-    #smoother.FeatureEdgeSmoothingOff()
-    #smoother.SetFeatureAngle(120.0)
-    #smoother.SetPassBand(.001)
-    #smoother.NonManifoldSmoothingOff()
-    #smoother.NormalizeCoordinatesOff()
-    #smoother.Update()
+    grid = vtk.vtkUnstructuredGrid()
+    grid.ShallowCopy(append_filter.GetOutput())
+    write_unstructured_grid_to_file(grid, "cosipy.vtu")
 
-    appendFilter = vtk.vtkAppendFilter()
-    appendFilter.AddInputData(delaunay.GetOutput())
-    appendFilter.Update()
 
-    unstructuredGrid = vtk.vtkUnstructuredGrid()
-    unstructuredGrid.ShallowCopy(appendFilter.GetOutput())
+def createDEM_v2(
+    file_path: str = "../../data/output/Zhadang_ERA5_20090101-20090110.nc",
+):
+    """Map dataset to DEM through indexing.
 
-    writer = vtk.vtkXMLUnstructuredGridWriter()
-    writer.SetFileName('cosipy.vtu')
-    writer.SetInputData(unstructuredGrid)
-    writer.Write()
+    .. note:: This method causes graphical glitches. Repeating
+        quad.GetPointIds for longitude gives a more accurate shape, but
+        the contours do not interpolate correctly so the colourmap
+        values are inaccurate.
+    """
 
-def createDEM_v2():
+    ds = xr.open_dataset(file_path)
+    check_2d(ds)
+    print("Creating DEM...\n")
+    points = create_points(data=ds)
 
-    ds = xr.open_dataset('../data/output/Peru_20160601-20180530_comp4.nc')
-    
-    points = vtk.vtkPoints()
+    print("Writing cells...\n")
     quad = vtk.vtkQuad()
     cells = vtk.vtkCellArray()
-    
-    numPoints = ds.south_north.size*ds.west_east.size
-   
-    print('Write points \n')
-    for i,j in product(ds.south_north.values,ds.west_east.values):
-            points.InsertNextPoint(ds.lat.isel(south_north=i,west_east=j), ds.lon.isel(south_north=i,west_east=j), ds.HGT.sel(south_north=i,west_east=j).values/6370000.0)
-    
-    print('Write cells \n')
-    for idx in range(points.GetNumberOfPoints()-ds.west_east.size):
-        if (idx%ds.west_east.size != 0):
-            quad.GetPointIds().SetId(0,idx)
-            quad.GetPointIds().SetId(1,idx+1)
-            quad.GetPointIds().SetId(2,idx+ds.west_east.size+1)
-            quad.GetPointIds().SetId(3,idx+ds.west_east.size)
+    for idx in range(points.GetNumberOfPoints() - ds.lat.size):
+        if idx % ds.lat.size != 0:
+            quad.GetPointIds().SetId(0, idx)
+            quad.GetPointIds().SetId(1, idx + 1)
+            quad.GetPointIds().SetId(2, idx + ds.lat.size + 1)
+            quad.GetPointIds().SetId(3, idx + ds.lat.size)
             cells.InsertNextCell(quad)
 
-    print('Create unstructured grid \n') 
-    grid = vtk.vtkUnstructuredGrid()
-    grid.SetPoints(points)
-    grid.SetCells(vtk.VTK_QUAD, cells)
-    
-    writer = vtk.vtkXMLUnstructuredGridWriter()
-    writer.SetFileName('cosipy.vtu')
-    writer.SetInputData(grid)
-    writer.Write()
-   
-def add_scalar(var, timestamp):
+    grid = create_unstructured_grid(points=points, cells=cells)
+    write_unstructured_grid_to_file(grid, "cosipy.vtu")
+
+
+def add_scalar(
+    file_path: str,
+    var: str,
+    timestamp: str,
+    mean: bool = False,
+    gltf: bool = False,
+):
+    """Selects data from array at specific time or as a daily mean.
+
+    .. todo:: support WRF coordinates (south_north/west_east)
+
+    Args:
+        file_path: Path to netcdf file.
+        var: Short name of variable in data.
+        timestamp: Time index of target data.
+        mean: If True, computes and selects the daily mean. Otherwise,
+            selects data at ``timestamp``. Default False.
+        gltf: If True, export as a 3D glTF object. Default False.
+
+    Returns:
+        Array selection at target time.
+    """
+
+    vtk_file = vtk.vtkXMLUnstructuredGridReader()
+    vtk_file.SetFileName("cosipy.vtu")
+    vtk_file.Update()
 
-    vtkFile = vtk.vtkXMLUnstructuredGridReader()
-    vtkFile.SetFileName('cosipy.vtu')
-    vtkFile.Update()
-    
     # Find cellId by coordinates
-    pointLocator =  vtk.vtkPointLocator()
-    pointLocator.SetDataSet(vtkFile.GetOutput())
-    pointLocator.BuildLocator()
-    
-    ds = xr.open_dataset('../data/output/Peru_20160601-20180530_comp4.nc')
-    ds = ds.sel(time=timestamp)
-    
-    ds_sub = ds[var].stack(x=['south_north','west_east']) 
-    ds_sub = ds_sub.dropna(dim='x')
-    lats = ds_sub.x.lat.values
-    lons = ds_sub.x.lon.values
+    point_locator = vtk.vtkPointLocator()
+    point_locator.SetDataSet(vtk_file.GetOutput())
+    point_locator.BuildLocator()
+
+    ds = xr.open_dataset(file_path)
+    ds = get_selection(array=ds, timestamp=timestamp, mean=mean)
+
+    latitude, longitude = get_coords(data=ds)
+    ds_sub = ds[var].stack(x=[latitude, longitude])
+    ds_sub = ds_sub.dropna(dim="x")
+    lats = ds_sub[latitude].values
+    lons = ds_sub[longitude].values
     data = ds_sub.values
-    print(lats)
 
-    numPoints = vtkFile.GetOutput().GetNumberOfPoints()
-    scalar = np.empty(numPoints)
+    num_points = vtk_file.GetOutput().GetNumberOfPoints()
+    scalar = np.empty(num_points)
     scalar[:] = np.nan
 
-    interpField = numpy_support.numpy_to_vtk(scalar)
-    interpField.SetName(var)
-    vtkFile.GetOutput().GetPointData().AddArray(interpField)
-    vtkFile.Update()
+    interp_field = numpy_support.numpy_to_vtk(scalar)
+    interp_field.SetName(var)
+    vtk_file.GetOutput().GetPointData().AddArray(interp_field)
+    vtk_file.Update()
 
-    print('Write points \n')
+    print("Writing points...\n")
     for i in np.arange(len(data)):
         # Get height
-        alt = ds.HGT.sel(lat=lats[i],lon=lons[i]).values/6370000.0
-        
-        pointId = vtk.mutable(0) 
-        Id = vtk.vtkIdList()
-        pointId = pointLocator.FindClosestPoint([lons[i],lats[i],alt])
-        vtkFile.GetOutput().GetPointData().GetArray(var).InsertTuple1(pointId,data[i])
+        alt = ds.HGT.sel(lat=lats[i], lon=lons[i]).values / 6370000.0
 
-    writer = vtk.vtkXMLUnstructuredGridWriter()
-    writer.SetFileName('cosipy.vtu')
-    writer.SetInputData(vtkFile.GetOutput())
-    writer.Write()
+        point_id = vtk.mutable(0.0)
+        point_id = point_locator.FindClosestPoint([lons[i], lats[i], alt])
+        if point_id < 0:
+            warnings.warn("tupleIndex from pointId < 0.")
+        else:
+            vtk_file.GetOutput().GetPointData().GetArray(var).InsertTuple1(
+                point_id, data[i]
+            )
 
-    #plotSurface(vtkFile)
+    write_unstructured_grid_to_file(vtk_file.GetOutput(), "cosipy.vtu")
 
-    
-def plotSurface(domain):
+    plotSurface(domain=vtk_file, variable=var, timestamp=timestamp, gltf=gltf)
 
-    print(domain)
-    domain.GetOutput().GetPointData().SetActiveScalars('TS')
-    
-    lut = vtk.vtkLookupTable()
-    lut.SetNumberOfColors(16)
-    lut.SetHueRange(2/3,0.0)
-    #lut.SetValueRange (270, 273);
-    lut.SetNanColor(1,1,1,0.5)
-    lut.SetNumberOfTableValues(3)
-    lut.SetNumberOfColors(16)
-    lut.Build()
 
-    # create a rendering window and renderer
-    ren = vtk.vtkRenderer()
-    renWin = vtk.vtkRenderWindow()
-    renWin.AddRenderer(ren)
-    WIDTH=1200
-    HEIGHT=1200
-    renWin.SetSize(WIDTH,HEIGHT)
-     
-    # create a renderwindowinteractor
-    joystickStyle = vtk.vtkInteractorStyleJoystickCamera()
-    iren = vtk.vtkRenderWindowInteractor()
-    iren.SetRenderWindow(renWin)
-    iren.SetInteractorStyle(joystickStyle)
+def convert_unstructured_to_polydata(
+    unstructured_grid: vtk.vtkUnstructuredGrid,
+) -> vtk.vtkPolyData:
+    """Converts unstructured grid to polydata.
 
-    # mapper
-    coneMapper = vtk.vtkDataSetMapper()
-    coneMapper.SetInputData(domain.GetOutput())
-    coneMapper.SetScalarModeToUsePointData()
-    coneMapper.ScalarVisibilityOn()
-    coneMapper.SetLookupTable(lut)
-    coneMapper.SetScalarRange(domain.GetOutput().GetPointData().GetArray("TS").GetRange())
+    More complex filters only accept polydata.
+    """
 
-    # actor
-    coneActor = vtk.vtkActor()
-    coneActor.SetMapper(coneMapper)
-    coneActor.SetScale(1,1,50)
-    coneActor.GetProperty().SetEdgeColor(0.7,0.7,0.7)
-    #coneActor.GetProperty().EdgeVisibilityOn()
-    coneActor.GetProperty().SetAmbient(0.0)
-    coneActor.GetProperty().SetSpecularPower(0)
-    coneActor.GetProperty().SetDiffuse(1)
-
-    cam = vtk.vtkCamera()
-    cam.SetPosition(10.89,46.83,0.07)
-    #cam.SetPosition(10.89,46.83,0.08)
-    cam.SetFocalPoint(10.82,46.83,0.02)
-    cam.SetViewUp(0,0,1)
-
-    # assign actor to the renderer
-    ren.SetBackground(1,1,1)
-    ren.AddActor(coneActor)
-    ren.SetActiveCamera(cam)
-    ren.ResetCamera()
+    geometry_filter = vtk.vtkGeometryFilter()
+    geometry_filter.SetInputData(unstructured_grid)
+    geometry_filter.Update()
+    polydata = geometry_filter.GetOutput()
+    msg = f"Converted to polydata with {polydata.GetNumberOfPoints()} points."
+    print(msg)
 
-    # enable user interface interactor
-    #iren.Initialize()
-    renWin.OffScreenRenderingOn()
-    renWin.Render()
-    #iren.Start()
-   
-    renderLarge = vtk.vtkRenderLargeImage()
-    renderLarge.SetInput(ren)
-    renderLarge.SetMagnification(5)
-
-#    w2if = vtk.vtkWindowToImageFilter()
-#    w2if.SetInputConnection(renderLarge.GetOutputPort())
-#    w2if.Update()
+    return polydata
 
-    writer = vtk.vtkPNGWriter()
-    writer.SetInputConnection(renderLarge.GetOutputPort())
-    writer.SetFileName('cosipy.png')
-    writer.Write()
 
+def set_filename(name: str, timestamp: str, fmt: str = "png") -> str:
+    """Creates a file name from timestamp and variable name.
 
-if __name__ == '__main__':
-    main()
+    Args:
+        name: Name of data variable.
+        timestamp: Time of data collection.
+        fmt: File format. Default "png".
 
+    Returns:
+        Output file name.
+    """
 
+    if not isinstance(timestamp, str):
+        img_id = timestamp.strftime("%Y%m%d")
+    else:
+        img_id = timestamp
+    img_id = re.sub(r"\W+", "", str(img_id))  # avoid illegal file names
+    img_id = f"{img_id}_{name}_vtk.{fmt}"
 
+    return img_id
 
 
+def plotSurface(
+    domain: vtk.vtkXMLUnstructuredGridReader,
+    variable: str,
+    timestamp: str = "",
+    contours: int = 10,
+    gltf: bool = False,
+) -> vtk.vtkRenderLargeImage:
+    """Plot surface with contoured data.
 
+    Args:
+        domain: Reader for glacier surface XML data.
+        variable: Short name of target variable.
+        timestamp: Time index of target data. Default empty string.
+        contours: Number of contours to plot. Default 10.
+        gltf: Export as 3D glTF object. Default False.
 
+    Returns:
+        Large image render of plotted data.
+    """
 
+    print(domain)
+    domain.GetOutput().GetPointData().SetActiveScalars(variable)
+    scalar_range = domain.GetOutput().GetPointData().GetArray(variable).GetRange()
+    num_contours = contours
+
+    lut = vtk.vtkLookupTable()
+    lut.SetNumberOfTableValues(num_contours + 1)
+    if variable.lower() == "hgt":
+        lut.SetHueRange(0.75, 0.1)  # blue: low altitude
+    else:
+        lut.SetHueRange(0.1, 0.75)  # matches cmap in plot_cosipy_profiles
+    lut.SetNanColor(1, 1, 1, 0.5)
+    lut.Build()
+
+    # mapper
+    cone_mapper = vtk.vtkDataSetMapper()
+    cone_mapper.SetInputData(
+        convert_unstructured_to_polydata(domain.GetOutput())  # for 3D export
+    )
+    cone_mapper.SetScalarModeToUsePointData()
+    cone_mapper.ScalarVisibilityOn()
+    cone_mapper.SetLookupTable(lut)
+    cone_mapper.SetScalarRange(scalar_range[0], scalar_range[1])
+    cone_mapper.SetInterpolateScalarsBeforeMapping(1)  # stops colour smudging
+    cone_mapper.Update()
+
+    # actor
+    cone_actor = vtk.vtkActor()
+    cone_actor.SetMapper(cone_mapper)
+    cone_actor.SetScale(1, 1, 50)
+    cone_actor.GetProperty().SetEdgeColor(0.7, 0.7, 0.7)
+    cone_actor.GetProperty().EdgeVisibilityOn()
+    cone_actor.GetProperty().SetAmbient(0.0)
+    cone_actor.GetProperty().SetSpecularPower(0)
+    cone_actor.GetProperty().SetDiffuse(1)
+    cone_actor.GetProperty().SetInterpolationToFlat()  # stop contour blending
+
+    # visualisation
+    camera = vtk.vtkCamera()
+    camera.SetPosition(0, 0, 1)  # x: west-east, y: south-north, z: top-down
+    camera.SetFocalPoint(0, 1, 0)
+    # cam.SetPosition(10.89, 46.83, 0.07)
+    # cam.SetPosition(10.89,46.83,0.08)
+    # cam.SetFocalPoint(10.82, 46.83, 0.02)
+    # cam.SetViewUp(0, 0, 1)
+
+    # create a rendering window and renderer
+    render = vtk.vtkRenderer()
+    render_window = vtk.vtkRenderWindow()
+    render_window.AddRenderer(render)
+    width = 1200
+    height = 1200
+    render_window.SetSize(width, height)
+
+    # create interactor for render window
+    joystick_style = vtk.vtkInteractorStyleJoystickCamera()
+    interactive_render = vtk.vtkRenderWindowInteractor()
+    interactive_render.SetRenderWindow(render_window)
+    interactive_render.SetInteractorStyle(joystick_style)
+    # assign actor to renderer
+    render.SetBackground(1, 1, 1)
+    render.AddActor(cone_actor)
+    render.SetActiveCamera(camera)
+    render.ResetCamera()
+
+    # enable user interface interactor
+    # interactive_render.Initialize()
+    render_window.OffScreenRenderingOn()
+    render_window.Render()
+    # interactive_render.Start()
+
+    render_large = vtk.vtkRenderLargeImage()
+    render_large.SetInput(render)
+    render_large.SetMagnification(5)
+
+    #    w2if = vtk.vtkWindowToImageFilter()
+    #    w2if.SetInputConnection(render_large.GetOutputPort())
+    #    w2if.Update()
+
+    write_render_to_image(render=render_large, variable=variable, timestamp=timestamp)
+
+    if gltf:
+        write_render_to_gltf(render=render, variable=variable, timestamp=timestamp)
+
+    return render_large
+
+
+def parse_arguments() -> argparse.Namespace:
+    """Parse user arguments.
+
+    Required arguments:
+        -i, --input <path>      Path to .nc file.
+        -t, --time <str>        Target date or timestamp.
+        -n, --name <str>        Name of variable to plot as contours.
+
+    Optional switches:
+        -h, --help              Show this help message and exit.
+        -m, --mean              Plot daily mean instead of timestep.
+                                    Default False.
+        -g, --gltf              Export plot as a 3D glTF object.
+                                    Default False.
+    """
+
+    tagline = "Plot 3D DEM with contours for a single variable."
+    parser = argparse.ArgumentParser(
+        prog="plot_cosipy_fields_vtk.py", description=tagline
+    )
+    # Required
+    parser.add_argument(
+        "-i",
+        "--input",
+        dest="input_file",
+        required=True,
+        default=None,
+        type=str,
+        metavar="<path>",
+        help="Path to .nc file",
+    )
+    parser.add_argument(
+        "-n",
+        "--name",
+        dest="name",
+        type=str,
+        default=None,
+        metavar="<str>",
+        help="Name of variable to plot as contours.",
+    )
+    parser.add_argument(
+        "-t",
+        "--time",
+        dest="timestamp",
+        type=str,
+        required=True,
+        default=None,
+        metavar="<str>",
+        help="Target date or timestamp",
+    )
+
+    # Optional switches
+    parser.add_argument(
+        "-m",
+        "--mean",
+        dest="mean",
+        action="store_true",
+        help="Plot daily mean instead of timestep",
+    )
+    parser.add_argument(
+        "-g",
+        "--gltf",
+        dest="gltf",
+        action="store_true",
+        help="Export plot as a 3D glTF object",
+    )
+
+    arguments = parser.parse_args()
+
+    return arguments
+
+
+if __name__ == "__main__":
+    main()
diff --git a/requirements.txt b/requirements.txt
index 390a8b5b..f6805679 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,10 +1,20 @@
-numpy
-pandas
-numba
-scipy
+numpy>=1.17.3
+pandas>2.0
+scipy>=1.1.0
 xarray>=0.18.0
-distributed<=2.5.2
-dask_jobqueue
-netcdf4
-coveralls
+netcdf4>=1.5.1.2
+distributed>=2022.3,<=2023.8; python_version < '3.11'
+distributed<=2023.2; python_version == '3.11'
+dask
+dask-jobqueue
+numba
 metpy
+matplotlib
+vtk; python_version < '3.11'
+richdem; python_version < '3.11'
+nco
+cdo
+cartopy
+coveralls
+codecov
+pytest


Test patch:
diff --git a/cosipy/tests/__init__.py b/cosipy/tests/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/cosipy/tests/conftest.py b/cosipy/tests/conftest.py
new file mode 100644
index 00000000..dfedd1a9
--- /dev/null
+++ b/cosipy/tests/conftest.py
@@ -0,0 +1,568 @@
+"""Provides shared fixtures and methods for tests.
+
+Use these to replace duplicated code.
+
+For generating objects within a test function's scope, call a fixture
+directly:
+
+    .. code-block:: python
+        
+        def test_foobar(self, conftest_mock_grid):
+            grid_object = conftest_mock_grid
+            grid_object.set_foo(foo=bar)
+            ...
+"""
+
+from types import ModuleType
+from typing import Any
+from unittest.mock import patch
+
+import matplotlib.pyplot as plt
+import numba
+import numpy as np
+import pandas as pd
+import pytest
+import xarray as xr
+
+from cosipy.cpkernel.grid import Grid
+
+
+# Function patches
+@pytest.fixture(scope="function", autouse=False)
+def conftest_mock_check_file_exists():
+    """Override checks when mocking files."""
+
+    patcher = patch("os.path.exists")
+    mock_exists = patcher.start()
+    mock_exists.return_value = True
+
+
+@pytest.fixture(scope="function", autouse=False)
+def conftest_disable_jit():
+    # numba.config.DISABLE_JIT = True
+    raise NotImplementedError(
+        "Disabling JIT for tests is not yet implemented."
+    )
+
+
+@pytest.fixture(scope="function", autouse=False)
+def conftest_mock_open_dataset(conftest_mock_xr_dataset):
+    """Override xr.open_dataset with mock dataset."""
+
+    patcher = patch("xarray.open_dataset")
+    mock_exists = patcher.start()
+    dataset = conftest_mock_xr_dataset.copy(deep=True)
+    mock_exists.return_value = dataset
+
+
+@pytest.fixture(scope="function", autouse=False)
+def conftest_hide_plot():
+    """Suppress plt.show(). Does not close plots."""
+
+    patcher = patch("matplotlib.pyplot.show")
+    mock_exists = patcher.start()
+    mock_exists.return_value = True
+
+
+@pytest.fixture(name="conftest_rng_seed", scope="function", autouse=False)
+def fixture_conftest_rng_seed() -> np.random.Generator:
+    """Set seed for random number generator to 444.
+
+    Returns:
+        Random number generator with seed=444.
+    """
+
+    random_generator = np.random.default_rng(seed=444)
+    assert isinstance(random_generator, np.random.Generator)
+
+    yield random_generator
+
+
+# Mock GRID data
+@pytest.fixture(
+    name="conftest_mock_grid_values", scope="function", autouse=False
+)
+def fixture_conftest_mock_grid_values() -> dict:
+    """Constructs the layer values used to generate Grid objects.
+
+    Returns:
+        Numba arrays for layers' heights, snowpack densities,
+        temperatures, and liquid water content.
+    """
+
+    layer_values = {}
+    layer_values["layer_heights"] = numba.float64([0.05, 0.1, 0.3, 0.5, 0.5])
+    layer_values["layer_densities"] = numba.float64([250, 250, 250, 917, 917])
+    layer_values["layer_temperatures"] = numba.float64(
+        [260, 270, 271, 271.5, 272]
+    )
+    layer_values["layer_liquid_water_content"] = numba.float64(
+        [0.0, 0.0, 0.0, 0.0, 0.0]
+    )
+
+    assert isinstance(layer_values, dict)
+    for array in layer_values.values():
+        assert isinstance(array, np.ndarray)
+        assert len(array) == 5
+
+    yield layer_values
+
+
+@pytest.fixture(name="conftest_mock_grid", scope="function", autouse=False)
+def fixture_conftest_mock_grid(conftest_mock_grid_values: dict) -> Grid:
+    """Constructs a Grid object.
+
+    .. note:: Use with caution, as this fixture assumes Grid objects are
+        correctly instantiated.
+
+    Returns:
+        Grid object with numba arrays for the layers' heights,
+        densities, temperatures, and liquid water content.
+    """
+
+    data = conftest_mock_grid_values.copy()
+    grid_object = Grid(
+        layer_heights=data["layer_heights"],
+        layer_densities=data["layer_densities"],
+        layer_temperatures=data["layer_temperatures"],
+        layer_liquid_water_content=data["layer_liquid_water_content"],
+    )
+    assert isinstance(grid_object, Grid)
+    assert grid_object.number_nodes == len(data["layer_heights"])
+
+    yield grid_object
+
+
+@pytest.fixture(name="conftest_mock_grid_ice", scope="function", autouse=False)
+def fixture_conftest_mock_grid_ice(conftest_mock_grid_values: dict) -> Grid:
+    """Constructs a Grid object for ice layers.
+
+    .. note:: Use with caution, as this fixture assumes Grid objects are
+        correctly instantiated.
+
+    Returns:
+        Grid object for ice layers with numba arrays for layer heights,
+        layer densities, layer temperatures, and layer liquid water
+        content.
+    """
+
+    data = conftest_mock_grid_values.copy()
+    grid_object = Grid(
+        layer_heights=data["layer_heights"][3:4],
+        layer_densities=data["layer_densities"][3:4],
+        layer_temperatures=data["layer_temperatures"][3:4],
+        layer_liquid_water_content=data["layer_liquid_water_content"][3:4],
+    )
+    assert isinstance(grid_object, Grid)
+
+    yield grid_object
+
+
+# Mock xarray Dataset
+@pytest.fixture(
+    name="conftest_mock_xr_dataset_dims", scope="function", autouse=False
+)
+def fixture_conftest_mock_xr_dataset_dims() -> dict:
+    """Yields dimensions for constructing an xr.Dataset."""
+
+    dimensions = {}
+    reference_time = pd.Timestamp("2009-01-01T12:00:00")
+    dimensions["time"] = pd.date_range(reference_time, periods=4, freq="6H")
+    dimensions["latitude"] = [30.460, 30.463, 30.469, 30.472]
+    dimensions["longitude"] = [90.621, 90.624, 90.627, 90.630, 90.633]
+    dimensions["name"] = ["time", "lat", "lon"]
+
+    yield dimensions
+
+
+@pytest.fixture(
+    name="conftest_mock_xr_dataset", scope="function", autouse=False
+)
+def fixture_conftest_mock_xr_dataset(
+    conftest_mock_xr_dataset_dims: dict, conftest_rng_seed: np.random.Generator
+) -> xr.Dataset:
+    """Constructs mock xarray Dataset of output .nc file.
+
+    Returns:
+        Dataset with data for elevation and surface data.
+    """
+
+    _ = conftest_rng_seed
+    dims = conftest_mock_xr_dataset_dims.copy()
+    lengths = [
+        len(dims["time"]),
+        len(dims["latitude"]),
+        len(dims["longitude"]),
+    ]
+    elevation = xr.Variable(
+        data=1000 + 10 * np.random.rand(lengths[1], lengths[2]),
+        dims=dims["name"][1:],
+        attrs={"long_name": "Elevation", "units": "m"},
+    )
+    temperature = xr.Variable(
+        data=15 + 8 * np.random.randn(lengths[0], lengths[1], lengths[2]),
+        dims=dims["name"],
+        attrs={"long_name": "Surface temperature", "units": "K"},
+    )
+
+    dataset = xr.Dataset(
+        data_vars=dict(HGT=elevation, TS=temperature),
+        coords=dict(
+            time=dims["time"],
+            lat=(["lat"], dims["latitude"]),
+            lon=(["lon"], dims["longitude"]),
+            reference_time=dims["time"][0],
+        ),
+        attrs=dict(
+            description="Weather related data.",
+            Full_fiels="True",  # match typo in io.py
+        ),
+    )
+    assert isinstance(dataset, xr.Dataset)
+
+    for key, length in zip(dims["name"], lengths):
+        assert dataset[key].shape == (length,)
+
+    assert "time" not in dataset.HGT.dims
+    assert dataset.HGT.shape == (lengths[1], lengths[2])
+    assert dataset.HGT.long_name == "Elevation"
+    for key in ["TS"]:  # in case we add more variables
+        assert "time" in dataset[key].dims
+        assert dataset[key].shape == (lengths[0], lengths[1], lengths[2])
+        assert dataset[key][0].shape == (lengths[1], lengths[2])
+
+    assert dataset.Full_fiels
+
+    yield dataset
+
+
+class TestBoilerplate:
+    """Provides boilerplate methods for serialising tests.
+
+    The class is instantiated via the `conftest_boilerplate` fixture.
+    The fixture is autoused, and can be called directly within a test::
+
+    ..code-block:: python
+
+        def test_foo(self, conftest_boilerplate):
+
+            foobar = [...]
+            conftest_boilerplate.bar(foobar)
+
+    Methods are arranged with their appropriate test::
+
+    .. code-block:: python
+
+        def foo(self, ...):
+            pass
+
+        def test_foo(self ...):
+            pass
+    """
+
+    def check_plot(self, plot_params: dict, title: str = ""):
+        """Check properties of figure/axis pairs.
+
+        Args:
+            plot_params: Contains matplotlib objects with supported keys:
+                "plot", "title", "x_label", "y_label".
+            title: Expected matplotlib figure title.
+        """
+
+        assert "plot" in plot_params
+        assert isinstance(plot_params["plot"][0], plt.Figure)
+        assert isinstance(plot_params["plot"][1], plt.Axes)
+        compare_ax = plot_params["plot"][1]
+        if "title" in plot_params:
+            compare_title = f"{plot_params['title']}{title}"
+            assert compare_ax.get_title("center") == compare_title
+        if "x_label" in plot_params:
+            assert compare_ax.xaxis.get_label_text() == plot_params["x_label"]
+        if "y_label" in plot_params:
+            assert compare_ax.yaxis.get_label_text() == plot_params["y_label"]
+
+    def test_check_plot(self):
+        """Validate tests for plot attributes."""
+
+        test_title = ", Inner Scope"
+        plt.close("all")
+        figure_strings = plt.figure()
+        axis_strings = plt.gca()
+        plt.title(f"Test Title{test_title}")
+        plt.xlabel("Test x-label")
+        plt.ylabel("Test y-label")
+
+        figure_none = plt.figure()
+        axis_none = plt.gca()
+
+        test_params = {
+            "strings": {
+                "plot": (figure_strings, axis_strings),
+                "title": "Test Title",
+                "x_label": "Test x-label",
+                "y_label": "Test y-label",
+            },
+            "none": {"plot": (figure_none, axis_none)},
+        }
+
+        for test_pair in test_params.values():
+            self.check_plot(plot_params=test_pair, title=test_title)
+        plt.close("all")
+
+    def set_timestamp(self, day: bool) -> str:
+        """Set timestamp string for a day or for a timestep."""
+
+        if day:
+            timestamp = "2009-01-01"
+        else:
+            timestamp = "2009-01-01T12:00:00"
+
+        return timestamp
+
+    def test_set_timestamp(self):
+        for arg_day in [True, False]:
+            compare_time = self.set_timestamp(day=arg_day)
+            assert isinstance(compare_time, str)
+            if arg_day:
+                assert compare_time == "2009-01-01"
+            else:
+                assert compare_time == "2009-01-01T12:00:00"
+
+    def set_rng_seed(self, seed: int = 444) -> np.random.Generator:
+        """Set seed for random number generator to 444.
+
+        Args:
+            seed: Seed value for generator.
+
+        Returns:
+            Random number generator with seed=444.
+        """
+
+        random_generator = np.random.default_rng(seed=seed)
+
+        return random_generator
+
+    def test_set_rng_seed(self):
+        rng_none = self.set_rng_seed()
+        assert isinstance(rng_none, np.random.Generator)
+        rng_none_state = rng_none.__getstate__()["state"]["state"]
+        rng_123 = self.set_rng_seed(seed=123)
+        assert isinstance(rng_123, np.random.Generator)
+        rng_123_state = rng_123.__getstate__()["state"]["state"]
+        rng_444 = self.set_rng_seed(seed=444)
+        assert isinstance(rng_444, np.random.Generator)
+        rng_444_state = rng_444.__getstate__()["state"]["state"]
+
+        assert all(
+            isinstance(state, int)
+            for state in [rng_none_state, rng_123_state, rng_444_state]
+        )
+        assert rng_none_state == rng_444_state
+        assert not rng_444_state == rng_123_state
+
+    def regenerate_grid_values(
+        self, grid: Grid, distribution: str = "zero"
+    ) -> Grid:
+        """Set liquid water content to match a distribution.
+
+        Args:
+            grid: Glacier data mesh.
+            distribution: Data distribution. Default "zero". Supports::
+
+            - random: Random uniform between 0.01 and 0.05.
+            - static: All liquid water content set to 0.1.
+            - decreasing: Follows `1 - (0.01 * node_index)`.
+            - increasing: Follows `0.01 * node_index`.
+            - zero: All liquid water content set to 0.
+
+        Returns:
+            Glacier data mesh with an updated distribution of liquid
+            water content.
+        """
+
+        if distribution == "random":
+            rng = self.set_rng_seed()
+            for idx in range(0, grid.number_nodes - 1):
+                grid.set_node_liquid_water_content(
+                    idx, rng.uniform(low=0.01, high=0.05)
+                )
+        elif distribution == "static":
+            for idx in range(0, grid.number_nodes - 1):
+                grid.set_node_liquid_water_content(idx, 0.1)
+        elif distribution == "decreasing":
+            for idx in range(0, grid.number_nodes - 1):
+                grid.set_node_liquid_water_content(idx, 1 - (0.01 * idx))
+        elif distribution == "increasing":
+            for idx in range(0, grid.number_nodes - 1):
+                grid.set_node_liquid_water_content(idx, 0.01 * idx)
+        elif distribution == "zero":
+            for idx in range(0, grid.number_nodes - 1):
+                grid.set_node_liquid_water_content(idx, 0)
+        else:
+            raise ValueError("Distribution not supported!")
+
+        return grid
+
+    def test_regenerate_grid_values(self, conftest_mock_grid):
+        grid = conftest_mock_grid
+        distribution_list = [
+            "static",
+            "decreasing",
+            "increasing",
+            "random",
+            "zero",
+        ]
+        for distribution in distribution_list:
+            grid = self.regenerate_grid_values(
+                grid=grid, distribution=distribution
+            )
+            assert isinstance(grid, Grid)
+            test_lwc = grid.get_liquid_water_content()
+            assert all(isinstance(i, float) for i in test_lwc)
+
+    def check_output(self, variable: Any, x_type: Any, x_value: Any) -> bool:
+        """Check a variable matches an expected type and value.
+
+        Args:
+            variable: Variable to check.
+            x_type: Expected variable type.
+            x_value: Expected variable value.
+
+        Returns:
+            True when all assertions pass.
+        """
+
+        assert isinstance(variable, x_type)
+        if np.issubdtype(type(variable), np.number):
+            assert np.isclose(variable, x_value)
+        else:
+            assert variable == x_value
+
+        return True
+
+    def test_check_output(self):
+        variable_list = [[1.0, float], ["test", str], [1, int], [True, bool]]
+
+        for pair in variable_list:
+            assert self.check_output(
+                variable=pair[0], x_type=pair[1], x_value=pair[0]
+            )
+        test_array = [0.0, 0.5, 0.6]
+        test_value = max(test_array)
+        assert test_value == 0.6
+        assert isinstance(test_value, float)
+        assert self.check_output(
+            variable=max(test_array), x_type=float, x_value=test_value
+        )
+
+    def assert_grid_profiles_equal(self, grid_01: Grid, grid_02: Grid) -> bool:
+        """Assert two Grid instances have equal grid profiles.
+
+        Args:
+            grid_01: Reference grid instance.
+            grid_02: Comparison grid instance.
+
+        Returns:
+            True when all assertions pass.
+        """
+        for grid in (grid_01, grid_02):
+            assert isinstance(grid, Grid)
+
+        assert grid_02 is not grid_01  # point to different instances
+        assert grid_02.number_nodes == grid_01.number_nodes
+        assert grid_02.get_number_layers() == grid_01.get_number_layers()
+        assert (
+            grid_02.get_number_snow_layers()
+            == grid_01.get_number_snow_layers()
+        )
+
+        self.check_output(
+            grid_02.get_total_height(), float, grid_01.get_total_height()
+        )
+        self.check_output(
+            grid_02.get_total_snowheight(),
+            float,
+            grid_01.get_total_snowheight(),
+        )
+        self.check_output(
+            sum(grid_02.get_ice_heights()),
+            float,
+            sum(grid_01.get_ice_heights()),
+        )
+
+        return True
+
+    def test_assert_grid_profiles_equal(self, conftest_mock_grid_values):
+        data = conftest_mock_grid_values.copy()
+
+        grid_01 = Grid(
+            layer_heights=data["layer_heights"],
+            layer_densities=data["layer_densities"],
+            layer_temperatures=data["layer_temperatures"],
+            layer_liquid_water_content=data["layer_liquid_water_content"],
+        )
+        grid_02 = Grid(
+            layer_heights=data["layer_heights"],
+            layer_densities=data["layer_densities"],
+            layer_temperatures=data["layer_temperatures"],
+            layer_liquid_water_content=data["layer_liquid_water_content"],
+        )
+        self.assert_grid_profiles_equal(grid_01, grid_02)
+
+    def patch_variable(
+        self,
+        monkeypatch: pytest.MonkeyPatch,
+        module: ModuleType,
+        new_params: dict,
+    ):
+        """Patch any variable in a module.
+
+        Patch the module where the variable is used, not where it's
+        defined. The patched variable only exists within the test
+        function's scope, so test parametrisation is still supported.
+
+        Example:
+            To patch constants used by `cpkernel.node.Node`:
+
+                .. code-block:: python
+
+                    patches = {"dt": 7200, "air_density": 1.0}
+                    conftest.boilerplate.patch_variable(
+                        monkeypatch,
+                        cosipy.cpkernel.node.constants,
+                        patches,
+                        )
+
+        Args:
+            monkeypatch: Monkeypatch instance.
+            module: Target module for patching.
+            new_params: Variable names as keys, desired patched values as values:
+
+                .. code-block:: python
+
+                    new_params = {"foo": 1, "bar": 2.0}
+        """
+
+        if not isinstance(new_params, dict):
+            note = "Pass dict with variable names and patched values as items."
+            raise TypeError(note)
+        for key in new_params:
+            monkeypatch.setattr(module, key, new_params[key])
+
+    @pytest.fixture(scope="class", autouse=False)
+    def test_boilerplate_integration(self):
+        """Integration test for boilerplate methods."""
+
+        self.test_check_plot()
+        self.test_set_timestamp()
+        self.test_set_rng_seed()
+        self.test_regenerate_grid_values()
+        self.test_check_output()
+        self.test_assert_grid_profiles_equal()
+
+
+@pytest.fixture(name="conftest_boilerplate", scope="function", autouse=False)
+def conftest_boilerplate():
+    """Yield class containing methods for common tests."""
+
+    yield TestBoilerplate()
diff --git a/cosipy/tests/test_GRID_getter_functions.py b/cosipy/tests/test_GRID_getter_functions.py
index fc7e9ea8..8f192a62 100644
--- a/cosipy/tests/test_GRID_getter_functions.py
+++ b/cosipy/tests/test_GRID_getter_functions.py
@@ -1,28 +1,203 @@
-from pytest import approx
-from cosipy.cpkernel.grid import *
-from COSIPY import start_logging
-
-layer_heights = [0.1, 0.2, 0.3, 0.5, 0.5]
-layer_densities = [250, 250, 250, 917, 917]
-layer_temperatures = [260, 270, 271, 271.5, 272]
-layer_liquid_water = [0.0, 0.0, 0.0, 0.0, 0.0]
-
-def test_grid_getter_functions():
-     GRID = Grid(layer_heights, layer_densities, layer_temperatures, layer_liquid_water)
-     assert GRID.get_height() == layer_heights
-     assert GRID.get_density() == approx(layer_densities, abs=1e-3)
-     assert GRID.get_temperature() == layer_temperatures
-     assert GRID.get_liquid_water_content() == layer_liquid_water
-     assert GRID.get_snow_heights() == layer_heights[0:3]
-     assert GRID.get_ice_heights() == layer_heights[3:5]
-     assert GRID.get_node_height(0) == layer_heights[0]
-     assert GRID.get_node_density(0) == approx(layer_densities[0], abs=1e-3)
-     assert GRID.get_node_temperature(0) == layer_temperatures[0]
-
-    # GRID.remove_melt_weq(0.01)
-    #########REMOVE_NODE
-    #number_nodes_before = GRID.get_number_layers()
-    #GRID.remove_node()
-    #assert GRID.get_number_nodes() == number_nodes_before -1
-    #assert np.nanmean(GRID.get_density()) == np.nanmean(layer_densities)
-    #assert GRID.get_density() == layer_densities
+import numpy as np
+import pytest
+from numba import float64
+
+import constants
+from cosipy.cpkernel.grid import Grid
+
+
+class TestGridSetup:
+    """Tests initialisation methods for Grid objects."""
+
+    def test_grid_init(self, conftest_boilerplate):
+        data = {
+            "layer_heights": [0.1, 0.2, 0.3],
+            "layer_densities": [200.0, 210.0, 220.0],
+            "layer_temperatures": [270.0, 260.0, 250.0],
+            "layer_liquid_water_content": [0.0, 0.0, 0.0],
+        }
+
+        test_grid = Grid(
+            layer_heights=float64(data["layer_heights"]),
+            layer_densities=float64(data["layer_densities"]),
+            layer_temperatures=float64(data["layer_temperatures"]),
+            layer_liquid_water_content=float64(
+                data["layer_liquid_water_content"]
+            ),
+        )
+        assert isinstance(test_grid, Grid)
+        conftest_boilerplate.check_output(
+            test_grid.number_nodes, int, len(data["layer_heights"])
+        )
+        assert test_grid.layer_ice_fraction is None
+        for i in range(test_grid.number_nodes):
+            assert isinstance(test_grid.get_node_ice_fraction(i), float)
+
+    def test_grid_init_grid(self, conftest_mock_grid):
+        test_grid = conftest_mock_grid
+        test_number_nodes = test_grid.number_nodes
+        test_grid.init_grid()
+
+        assert test_grid.number_nodes == test_number_nodes
+        assert test_grid.layer_ice_fraction is None
+        for i in range(test_grid.number_nodes):
+            assert isinstance(test_grid.get_node_ice_fraction(i), float)
+
+
+class TestGridGetter:
+    """Tests get methods for Grid objects.
+
+    ..
+        Pytest documentation recommends `np.allclose` instead of
+        `pytest.approx`.
+
+    Attributes:
+        data (dict[float64]): Dummy grid data.
+    """
+
+    data = {
+        "layer_heights": float64([0.1, 0.2, 0.3, 0.5, 0.5]),
+        "layer_densities": float64([250, 250, 250, 917, 917]),
+        "layer_temperatures": float64([260, 270, 271, 271.5, 272]),
+        "layer_liquid_water_content": float64([0.0, 0.0, 0.0, 0.0, 0.0]),
+    }
+
+    def create_grid(self):
+        grid_object = Grid(
+            layer_heights=self.data["layer_heights"],
+            layer_densities=self.data["layer_densities"],
+            layer_temperatures=self.data["layer_temperatures"],
+            layer_liquid_water_content=self.data["layer_liquid_water_content"],
+        )
+        return grid_object
+
+    @pytest.fixture(name="grid", autouse=False, scope="function")
+    def fixture_grid(self):
+        return self.create_grid()
+
+    def get_ice_fraction(self, ice_fraction, density):
+        if ice_fraction is None:
+            a = (
+                density
+                - (1 - (density / constants.ice_density))
+                * constants.air_density
+            )
+            ice_fraction = a / constants.ice_density
+        else:
+            ice_fraction = ice_fraction
+        return ice_fraction
+
+    @pytest.mark.parametrize("arg_ice_fraction", [0.1, None])
+    def test_get_ice_fraction(self, arg_ice_fraction, conftest_boilerplate):
+        test_density = 300
+        if arg_ice_fraction is None:
+            a = (
+                test_density
+                - (1 - (test_density / constants.ice_density))
+                * constants.air_density
+            )
+            test_ice = a / constants.ice_density
+        else:
+            test_ice = arg_ice_fraction
+        compare_ice = self.get_ice_fraction(arg_ice_fraction, test_density)
+        assert conftest_boilerplate.check_output(compare_ice, float, test_ice)
+
+    def test_create_grid(self):
+        test_grid = self.create_grid()
+        assert isinstance(test_grid, Grid)
+        assert test_grid.number_nodes == len(self.data["layer_heights"])
+
+    def test_grid_get_height(self, grid, conftest_boilerplate):
+        assert np.allclose(grid.get_height(), self.data["layer_heights"])
+        for i in range(grid.number_nodes):
+            conftest_boilerplate.check_output(
+                grid.get_node_height(i), float, self.data["layer_heights"][i]
+            )
+
+    def test_grid_get_temperature(self, grid, conftest_boilerplate):
+        assert np.allclose(
+            grid.get_temperature(), self.data["layer_temperatures"]
+        )
+        for i in range(grid.number_nodes):
+            conftest_boilerplate.check_output(
+                grid.get_node_temperature(i),
+                float,
+                self.data["layer_temperatures"][i],
+            )
+
+    def test_grid_get_liquid_water_content(self, grid, conftest_boilerplate):
+        assert np.allclose(
+            grid.get_liquid_water_content(),
+            self.data["layer_liquid_water_content"],
+        )
+        for i in range(grid.number_nodes):
+            conftest_boilerplate.check_output(
+                grid.get_node_liquid_water_content(i),
+                float,
+                self.data["layer_liquid_water_content"][i],
+            )
+
+    def test_grid_get_density(self, grid, conftest_boilerplate):
+        assert np.allclose(grid.get_density(), self.data["layer_densities"])
+        for i in range(grid.number_nodes):
+            conftest_boilerplate.check_output(
+                grid.get_node_density(i),
+                float,
+                self.data["layer_densities"][i],
+            )
+
+    def test_grid_get_ice_fraction(self, grid, conftest_boilerplate):
+        ice_fractions = [
+            self.get_ice_fraction(None, density)
+            for density in self.data["layer_densities"]
+        ]
+        assert np.allclose(grid.get_ice_fraction(), ice_fractions)
+
+        for i in range(grid.number_nodes):
+            conftest_boilerplate.check_output(
+                grid.get_node_ice_fraction(i), float, ice_fractions[i]
+            )
+
+    def test_grid_get_refreeze(self, grid, conftest_boilerplate):
+        refrozen = [0.0 for i in range(grid.number_nodes)]
+        assert np.allclose(grid.get_refreeze(), refrozen)
+
+        for i in range(grid.number_nodes):
+            conftest_boilerplate.check_output(
+                grid.get_node_refreeze(i), float, refrozen[i]
+            )
+
+    def test_grid_get_snow_ice_heights(
+        self, conftest_mock_grid_values, conftest_mock_grid
+    ):
+        data = conftest_mock_grid_values.copy()
+        GRID = conftest_mock_grid
+
+        assert np.allclose(GRID.get_snow_heights(), data["layer_heights"][0:3])
+        assert np.allclose(GRID.get_ice_heights(), data["layer_heights"][3:5])
+        assert np.allclose(GRID.get_node_height(0), data["layer_heights"][0])
+
+    def test_grid_getter_functions(
+        self, conftest_mock_grid_values, conftest_mock_grid
+    ):
+        data = conftest_mock_grid_values.copy()
+        GRID = conftest_mock_grid
+
+        # pytest documentation recommends np.allclose instead of pytest.approx
+        assert np.allclose(GRID.get_height(), data["layer_heights"])
+        assert np.allclose(
+            GRID.get_density(), data["layer_densities"], atol=1e-3
+        )
+        assert np.allclose(GRID.get_temperature(), data["layer_temperatures"])
+        assert np.allclose(
+            GRID.get_liquid_water_content(), data["layer_liquid_water_content"]
+        )
+        assert np.allclose(GRID.get_snow_heights(), data["layer_heights"][0:3])
+        assert np.allclose(GRID.get_ice_heights(), data["layer_heights"][3:5])
+        assert np.allclose(GRID.get_node_height(0), data["layer_heights"][0])
+        assert np.allclose(
+            GRID.get_node_density(0), data["layer_densities"][0], atol=1e-3
+        )
+        assert np.allclose(
+            GRID.get_node_temperature(0), data["layer_temperatures"][0]
+        )
diff --git a/cosipy/tests/test_GRID_update_functions.py b/cosipy/tests/test_GRID_update_functions.py
index 912f96c9..ab322547 100644
--- a/cosipy/tests/test_GRID_update_functions.py
+++ b/cosipy/tests/test_GRID_update_functions.py
@@ -1,33 +1,211 @@
-from pytest import approx
-from cosipy.cpkernel.grid import *
+import numpy as np
+import pytest
+
+# import constants
 from COSIPY import start_logging
+from cosipy.cpkernel.node import Node
+
+
+class TestGridUpdate:
+    """Tests update methods for Grid objects."""
+
+    def test_grid_set_node_height(
+        self, conftest_mock_grid, conftest_boilerplate
+    ):
+        GRID = conftest_mock_grid
+        test_height = GRID.get_node_height(0)
+        GRID.set_node_height(0, test_height + 0.5)
+        compare_height = GRID.get_node_height(0)
+        conftest_boilerplate.check_output(
+            compare_height, float, test_height + 0.5
+        )
+
+    def test_grid_set_node_temperature(
+        self, conftest_mock_grid, conftest_boilerplate
+    ):
+        GRID = conftest_mock_grid
+        test_temperature = GRID.get_node_temperature(0)
+        GRID.set_node_temperature(0, test_temperature + 0.5)
+        compare_temperature = GRID.get_node_temperature(0)
+        conftest_boilerplate.check_output(
+            compare_temperature, float, test_temperature + 0.5
+        )
+
+    def test_grid_set_node_ice_fraction(
+        self, conftest_mock_grid, conftest_boilerplate
+    ):
+        GRID = conftest_mock_grid
+        test_fraction = GRID.get_node_ice_fraction(0)
+        GRID.set_node_ice_fraction(0, test_fraction + 0.5)
+        compare_fraction = GRID.get_node_ice_fraction(0)
+        conftest_boilerplate.check_output(
+            compare_fraction, float, test_fraction + 0.5
+        )
+
+    def test_grid_set_node_liquid_water_content(
+        self, conftest_mock_grid, conftest_boilerplate
+    ):
+        GRID = conftest_mock_grid
+        test_lwc = GRID.get_node_liquid_water_content(0)
+        GRID.set_node_liquid_water_content(0, test_lwc + 0.5)
+        compare_lwc = GRID.get_node_liquid_water_content(0)
+        conftest_boilerplate.check_output(compare_lwc, float, test_lwc + 0.5)
+
+    def test_grid_set_node_refreeze(
+        self, conftest_mock_grid, conftest_boilerplate
+    ):
+        GRID = conftest_mock_grid
+        test_refreeze = GRID.get_node_liquid_water_content(0)
+        GRID.set_node_liquid_water_content(0, test_refreeze + 0.5)
+        compare_refreeze = GRID.get_node_liquid_water_content(0)
+        conftest_boilerplate.check_output(
+            compare_refreeze, float, test_refreeze + 0.5
+        )
+
+    def test_grid_update_functions(self, conftest_mock_grid):
+        GRID = conftest_mock_grid
+        GRID.set_node_liquid_water_content(0, 0.04)
+        GRID.set_node_liquid_water_content(1, 0.03)
+        GRID.set_node_liquid_water_content(2, 0.03)
+        GRID.set_node_liquid_water_content(3, 0.02)
+        GRID.set_node_liquid_water_content(4, 0.01)
+
+        SWE_before = np.array(GRID.get_height()) / np.array(GRID.get_density())
+        SWE_before_sum = np.nansum(SWE_before)
+
+        GRID.update_grid()
+        SWE_after = np.array(GRID.get_height()) / np.array(GRID.get_density())
+        SWE_after_sum = np.nansum(SWE_after)
+
+        GRID.adaptive_profile()
+        SWE_after_adaptive = np.array(GRID.get_height()) / np.array(
+            GRID.get_density()
+        )
+        SWE_after_adaptive_sum = np.nansum(SWE_after_adaptive)
+
+        assert np.allclose(SWE_before_sum, SWE_after_sum, atol=1e-3)
+        assert np.allclose(SWE_after_sum, SWE_after_adaptive_sum, atol=1e-3)
+
+
+class TestGridInteractions:
+    """Tests remeshing and interactions between layers."""
+
+    @pytest.mark.parametrize("arg_height", [0.05, 0.1, 0.5])
+    @pytest.mark.parametrize("arg_temperature", [273.16, 270.16, 280.0])
+    @pytest.mark.parametrize("arg_lwc", [0.0, 0.5])
+    def test_add_fresh_snow(
+        self,
+        conftest_mock_grid,
+        conftest_boilerplate,
+        arg_height,
+        arg_temperature,
+        arg_lwc,
+    ):
+        """Add fresh snow layer."""
+
+        test_grid = conftest_mock_grid
+        test_number_nodes = test_grid.number_nodes
+        test_snow = test_grid.get_fresh_snow_props()
+        assert isinstance(test_snow, tuple)
+        assert all(isinstance(parameter, float) for parameter in test_snow)
+        assert test_snow[0] == 0
+
+        test_grid.add_fresh_snow(arg_height, 250.0, arg_temperature, arg_lwc)
+        assert test_grid.number_nodes == test_number_nodes + 1
+        assert isinstance(test_grid.grid[0], Node)
+
+        fresh_snow = test_grid.get_fresh_snow_props()
+        assert isinstance(fresh_snow, tuple)
+        assert all(isinstance(parameter, float) for parameter in fresh_snow)
+        assert conftest_boilerplate.check_output(
+            fresh_snow[0], float, arg_height
+        )
+        assert not np.isclose(fresh_snow[0], test_snow[0])
+
+        compare_node = test_grid.grid[0]
+        conftest_boilerplate.check_output(
+            compare_node.height, float, arg_height
+        )
+        conftest_boilerplate.check_output(
+            compare_node.temperature, float, arg_temperature
+        )
+        conftest_boilerplate.check_output(
+            compare_node.liquid_water_content, float, arg_lwc
+        )
+
+    @pytest.mark.parametrize("arg_idx", [None, [-1], [0, 1, 2]])
+    def test_grid_remove_node(
+        self, conftest_mock_grid_values, conftest_mock_grid, arg_idx
+    ):
+        """Remove node from grid with or without indices."""
+
+        data = conftest_mock_grid_values.copy()
+        GRID = conftest_mock_grid
+        if not arg_idx:
+            indices = [0]
+        else:
+            indices = arg_idx
+        assert isinstance(indices, list)
+
+        GRID.remove_melt_weq(0.01)
+        number_nodes_before = GRID.get_number_layers()
+        GRID.remove_node(arg_idx)  # Remove node
+
+        assert GRID.get_number_layers() == number_nodes_before - len(indices)
+        assert np.isclose(  # matches new density
+            np.nanmean(GRID.get_density()),
+            np.nanmean(np.delete(data["layer_densities"], indices)),
+        )
+
+
+# class TestGridUpdate:
+#     """Tests update methods for Grid objects."""
+
+#     def test_grid_update_functions(self, conftest_mock_grid):
+#         GRID = conftest_mock_grid
+#         GRID.set_node_liquid_water_content(0, 0.04)
+#         GRID.set_node_liquid_water_content(1, 0.03)
+#         GRID.set_node_liquid_water_content(2, 0.03)
+#         GRID.set_node_liquid_water_content(3, 0.02)
+#         GRID.set_node_liquid_water_content(4, 0.01)
+
+#         SWE_before = np.array(GRID.get_height()) / np.array(GRID.get_density())
+#         SWE_before_sum = np.nansum(SWE_before)
+
+#         GRID.update_grid()
+#         SWE_after = np.array(GRID.get_height()) / np.array(GRID.get_density())
+#         SWE_after_sum = np.nansum(SWE_after)
 
-layer_heights = [0.1, 0.2, 0.3, 0.5, 0.5]
-layer_densities = [250, 250, 250, 917, 917]
-layer_temperatures = [260, 270, 271, 271.5, 272]
-layer_liquid_water = [0.0, 0.0, 0.0, 0.0, 0.0]
+#         GRID.adaptive_profile()
+#         SWE_after_adaptive = np.array(GRID.get_height()) / np.array(
+#             GRID.get_density()
+#         )
+#         SWE_after_adaptive_sum = np.nansum(SWE_after_adaptive)
 
-def test_grid_getter_functions():
-     GRID = Grid(layer_heights, layer_densities, layer_temperatures, layer_liquid_water)
-     GRID.set_node_liquid_water_content(0, 0.04)
-     GRID.set_node_liquid_water_content(1, 0.03)
-     GRID.set_node_liquid_water_content(2, 0.03)
-     GRID.set_node_liquid_water_content(3, 0.02)
-     GRID.set_node_liquid_water_content(4, 0.01)
+#         assert np.allclose(SWE_before_sum, SWE_after_sum, atol=1e-3)
+#         assert np.allclose(SWE_after_sum, SWE_after_adaptive_sum, atol=1e-3)
 
-     SWE_before = np.array(GRID.get_height()) / np.array(GRID.get_density())
-     SWE_before_sum = np.nansum(SWE_before)
+#     @pytest.mark.parametrize("arg_height", [0.05, 0.1, 0.5])
+#     @pytest.mark.parametrize("arg_temperature", [273.16, 270.16, 280.0])
+#     def test_add_fresh_snow(
+#         self, conftest_mock_grid, arg_height, arg_temperature
+#     ):
+#         """Add fresh snow layer."""
 
-     GRID.update_grid()
-     SWE_after = np.array(GRID.get_height()) / np.array(GRID.get_density())
-     SWE_after_sum = np.nansum(SWE_after)
+#         grid = conftest_mock_grid
+#         assert isinstance(grid, Grid)
 
-     GRID.adaptive_profile()
-     SWE_after_adaptive = np.array(GRID.get_height()) / np.array(GRID.get_density())
-     SWE_after_adaptive_sum = np.nansum(SWE_after_adaptive)
+#         snow = grid.get_fresh_snow_props()
+#         assert isinstance(snow, tuple)
+#         assert all(isinstance(parameter, float) for parameter in snow)
+#         assert snow[0] == 0
 
-     #assert SWE_before_sum == SWE_after_sum
-     #assert SWE_after_sum == SWE_after_adaptive_sum
+#         # grid.add_fresh_snow(height=arg_height, density=250, temperature=270.15, liquid_water_content=0.0)
+#         grid.add_fresh_snow(arg_height, 250.0, arg_temperature, 0.0)
+#         assert isinstance(grid, Grid)
 
-     assert SWE_before_sum == approx(SWE_after_sum, abs=1e-3)
-     assert SWE_after_sum == approx(SWE_after_adaptive_sum, abs=1e-3)
\ No newline at end of file
+#         fresh_snow = grid.get_fresh_snow_props()
+#         assert isinstance(fresh_snow, tuple)
+#         assert all(isinstance(parameter, float) for parameter in fresh_snow)
+#         assert np.isclose(fresh_snow[0], arg_height)
+#         assert not np.isclose(fresh_snow[0], snow[0])
diff --git a/cosipy/tests/test_NODE_getter_functions.py b/cosipy/tests/test_NODE_getter_functions.py
index deb409e8..e3b32293 100644
--- a/cosipy/tests/test_NODE_getter_functions.py
+++ b/cosipy/tests/test_NODE_getter_functions.py
@@ -1,57 +1,227 @@
-from cosipy.cpkernel.node import *
-
-height = 0.1
-density = 200.
-temperature = 270
-lwc = 0.2
-ice_fraction = 0.4
-
-test_node = Node(height, density, temperature, lwc)
-test_node.set_layer_ice_fraction(ice_fraction)
-
-def calculate_irreducible_water_content(current_ice_fraction):
-    if (current_ice_fraction <= 0.23):
-        theta_e = 0.0264 + 0.0099 * ((1 - current_ice_fraction) / current_ice_fraction)
-    elif (current_ice_fraction > 0.23) & (current_ice_fraction <= 0.812):
-        theta_e = 0.08 - 0.1023 * (current_ice_fraction - 0.03)
-    else:
-        theta_e = 0.0
-    return theta_e
-
-def test_node_getter_functions():
-    assert test_node.get_layer_height() == height
-
-    assert test_node.get_layer_temperature() == temperature
-
-    assert test_node.get_layer_ice_fraction() == ice_fraction
-
-    assert test_node.get_layer_refreeze() == 0.0
-
-    assert test_node.get_layer_density() == test_node.get_layer_ice_fraction()*ice_density + \
-           test_node.get_layer_liquid_water_content()*water_density + test_node.get_layer_air_porosity()*air_density
-
-    assert test_node.get_layer_air_porosity() == 1 - lwc - ice_fraction
-
-    assert test_node.get_layer_specific_heat() == (1 - lwc - ice_fraction) * spec_heat_air + ice_fraction * spec_heat_ice \
-           + lwc * spec_heat_water
-
-    assert test_node.get_layer_liquid_water_content() == lwc
-
-    assert test_node.get_layer_irreducible_water_content() == calculate_irreducible_water_content(test_node.get_layer_ice_fraction())
-
-    assert test_node.get_layer_cold_content() == - test_node.get_layer_specific_heat() * test_node.get_layer_density() \
-           * height * (temperature - zero_temperature)
-
-    assert test_node.get_layer_porosity() == 1 - test_node.get_layer_ice_fraction() - test_node.get_layer_liquid_water_content()
-
-    assert test_node.get_layer_thermal_conductivity() == ice_fraction * k_i + test_node.get_layer_porosity() * k_a + lwc * k_w
-
-    assert test_node.get_layer_thermal_diffusivity() == test_node.get_layer_thermal_conductivity() \
-           / (test_node.get_layer_density() * test_node.get_layer_specific_heat())
-
-def test_node_getter_functions_other_cases():
-    test_node.set_layer_ice_fraction(0.1)
-    assert test_node.get_layer_irreducible_water_content() == calculate_irreducible_water_content(test_node.get_layer_ice_fraction())
-
-    test_node.set_layer_ice_fraction(0.9)
-    assert test_node.get_layer_irreducible_water_content() == calculate_irreducible_water_content(test_node.get_layer_ice_fraction())
+import pytest
+
+import constants
+from cosipy.cpkernel.node import Node
+
+
+class TestNodeGetter:
+    """Tests get methods for Node objects.
+
+    Attributes:
+        height (float): Layer height [:math:`m`]
+        snow_density (float): Snow density [:math:`kg~m^{-3}`]
+        temperature (int): Layer temperature [:math:`K`]
+        lwc (float): Liquid water content [:math:`m~w.e.`]
+        ice_fraction (float): Volumetric ice fraction [-]
+    """
+
+    height = 0.1
+    snow_density = 200.0
+    temperature = 270.0
+    lwc = 0.2
+    ice_fraction = 0.4
+
+    def test_node_init(self, conftest_boilerplate):
+        """Inherit methods from parent."""
+
+        test_node = Node(
+            height=self.height,
+            snow_density=self.snow_density,
+            temperature=self.temperature,
+            liquid_water_content=self.lwc,
+            ice_fraction=self.ice_fraction,
+        )
+        assert test_node
+        assert isinstance(test_node, Node)
+
+        conftest_boilerplate.check_output(
+            variable=test_node.temperature,
+            x_type=float,
+            x_value=self.temperature,
+        )
+        conftest_boilerplate.check_output(test_node.height, float, self.height)
+        conftest_boilerplate.check_output(
+            test_node.liquid_water_content, float, self.lwc
+        )
+        conftest_boilerplate.check_output(test_node.refreeze, float, 0.0)
+        # conftest_boilerplate.check_output(
+        #     test_node.snow_density, float, self.snow_density
+        # )
+
+    def create_node(
+        self,
+        height: float = height,
+        snow_density: float = snow_density,
+        temperature: float = temperature,
+        lwc: float = lwc,
+        ice_fraction: float = ice_fraction,
+    ) -> Node:
+        """Instantiate a Node."""
+
+        node = Node(
+            height=height,
+            snow_density=snow_density,
+            temperature=temperature,
+            liquid_water_content=lwc,
+            ice_fraction=ice_fraction,
+        )
+        assert isinstance(node, Node)
+
+        return node
+
+    def test_create_node(self):
+        node = self.create_node()
+        assert isinstance(node, Node)
+
+    def calculate_irreducible_water_content(
+        self, current_ice_fraction: float
+    ) -> float:
+        """Calculate irreducible water content."""
+        if current_ice_fraction <= 0.23:
+            theta_e = 0.0264 + 0.0099 * (
+                (1 - current_ice_fraction) / current_ice_fraction
+            )
+        elif (current_ice_fraction > 0.23) & (current_ice_fraction <= 0.812):
+            theta_e = 0.08 - 0.1023 * (current_ice_fraction - 0.03)
+        else:
+            theta_e = 0.0
+
+        return theta_e
+
+    @pytest.mark.parametrize("arg_ice_fraction", [0.2, 0.5, 0.9])
+    def test_calculate_irreducible_water_content(self, arg_ice_fraction):
+        theta_e = self.calculate_irreducible_water_content(arg_ice_fraction)
+        assert isinstance(theta_e, float)
+
+    @pytest.fixture(name="node", autouse=False, scope="function")
+    def fixture_node(self):
+        return self.create_node()
+
+    def test_node_get_layer_height(self, node, conftest_boilerplate):
+        assert conftest_boilerplate.check_output(
+            node.get_layer_height(), float, self.height
+        )
+
+    def test_node_get_layer_temperature(self, node, conftest_boilerplate):
+        assert conftest_boilerplate.check_output(
+            node.get_layer_temperature(), float, self.temperature
+        )
+
+    def test_node_get_layer_liquid_water_content(
+        self, node, conftest_boilerplate
+    ):
+        assert conftest_boilerplate.check_output(
+            node.get_layer_liquid_water_content(), float, self.lwc
+        )
+
+    def test_node_get_layer_refreeze(self, node, conftest_boilerplate):
+        assert conftest_boilerplate.check_output(
+            node.get_layer_refreeze(), float, 0.0
+        )
+
+    @pytest.mark.parametrize("arg_ice_fraction", [0.0, 0.1, 0.9, None])
+    def test_node_get_layer_ice_fraction(
+        self, conftest_boilerplate, arg_ice_fraction
+    ):
+        node = self.create_node(ice_fraction=arg_ice_fraction)
+        if arg_ice_fraction is None:
+            test_ice_fraction = (
+                self.snow_density
+                - (1 - (self.snow_density / constants.ice_density))
+                * constants.air_density
+            ) / constants.ice_density
+        else:
+            test_ice_fraction = arg_ice_fraction
+        compare_ice_fraction = node.get_layer_ice_fraction()
+        conftest_boilerplate.check_output(
+            compare_ice_fraction, float, test_ice_fraction
+        )
+        conftest_boilerplate.check_output(
+            compare_ice_fraction, float, node.ice_fraction
+        )
+
+    def test_node_get_layer_air_porosity(self, node, conftest_boilerplate):
+        test_porosity = 1 - self.lwc - self.ice_fraction
+        conftest_boilerplate.check_output(
+            node.get_layer_air_porosity(), float, test_porosity
+        )
+
+    def test_node_get_layer_density(self, node, conftest_boilerplate):
+        test_density = (
+            self.ice_fraction * constants.ice_density
+            + self.lwc * constants.water_density
+            + node.get_layer_air_porosity() * constants.air_density
+        )
+        assert conftest_boilerplate.check_output(
+            node.get_layer_density(), float, test_density
+        )
+
+    def test_node_get_layer_porosity(self, node, conftest_boilerplate):
+        test_porosity = 1 - self.lwc - self.ice_fraction
+        conftest_boilerplate.check_output(
+            node.get_layer_porosity(), float, test_porosity
+        )
+
+    def test_node_get_layer_specific_heat(self, node, conftest_boilerplate):
+        test_specific_heat = (
+            (1 - self.lwc - self.ice_fraction) * constants.spec_heat_air
+            + self.ice_fraction * constants.spec_heat_ice
+            + self.lwc * constants.spec_heat_water
+        )
+        conftest_boilerplate.check_output(
+            node.get_layer_specific_heat(), float, test_specific_heat
+        )
+
+    def test_node_get_layer_cold_content(self, node, conftest_boilerplate):
+        test_cold_content = (
+            -node.get_layer_specific_heat()
+            * node.get_layer_density()
+            * self.height
+            * (self.temperature - constants.zero_temperature)
+        )
+        conftest_boilerplate.check_output(
+            node.get_layer_cold_content(), float, test_cold_content
+        )
+
+    def test_node_get_layer_thermal_conductivity(
+        self, node, conftest_boilerplate
+    ):
+        test_thermal_conductivity = (
+            self.ice_fraction * constants.k_i
+            + node.get_layer_porosity() * constants.k_a
+            + self.lwc * constants.k_w
+        )
+        conftest_boilerplate.check_output(
+            node.get_layer_thermal_conductivity(),
+            float,
+            test_thermal_conductivity,
+        )
+
+    def test_node_get_layer_thermal_diffusivity(
+        self, node, conftest_boilerplate
+    ):
+        test_thermal_diffusivity = node.get_layer_thermal_conductivity() / (
+            node.get_layer_density() * node.get_layer_specific_heat()
+        )
+        conftest_boilerplate.check_output(
+            node.get_layer_thermal_diffusivity(),
+            float,
+            test_thermal_diffusivity,
+        )
+
+    @pytest.mark.parametrize("arg_ice_fraction", [0.1, 0.7, 0.9])
+    def test_node_get_layer_irreducible_water_content(
+        self, conftest_boilerplate, arg_ice_fraction
+    ):
+        node = self.create_node(ice_fraction=arg_ice_fraction)
+
+        test_irreducible_water_content = (
+            self.calculate_irreducible_water_content(
+                node.get_layer_ice_fraction()
+            )
+        )
+        conftest_boilerplate.check_output(
+            node.get_layer_irreducible_water_content(),
+            float,
+            test_irreducible_water_content,
+        )
diff --git a/cosipy/tests/test_NODE_setter_functions.py b/cosipy/tests/test_NODE_setter_functions.py
index a900c855..ace179ed 100644
--- a/cosipy/tests/test_NODE_setter_functions.py
+++ b/cosipy/tests/test_NODE_setter_functions.py
@@ -1,30 +1,98 @@
-from cosipy.cpkernel.node import *
-height = 0.1
-density = 200.
-temperature = 270
-lwc = 0.3
-ice_fraction = 0.4
-test_node = Node(height, density, temperature, lwc, ice_fraction)
+# import constants
+from cosipy.cpkernel.node import Node
 
-new_height = 1.
-new_temperature = 260.
-new_liquidwater_content = 0.5
-new_ice_fraction = 0.9
-new_refreeze = 0.01
 
-def test_node_setter_functions():
+class TestNodeSetter:
+    """Tests set methods for Node objects.
 
-    test_node.set_layer_height(new_height)
-    assert test_node.get_layer_height() == new_height
+    Attributes:
+        height (float): Layer height [:math:`m`].
+        density (float): Snow density [:math:`kg~m^{-3}`].
+        temperature (int): Layer temperature [:math:`K`].
+        lwc (float): Liquid water content [:math:`m~w.e.`].
+        ice_fraction (float): Volumetric ice fraction [-].
+        new_height (float): New layer height [:math:`m`].
+        new_temperature (int): New layer temperature [:math:`K`].
+        new_lwc (float): New liquid water content [:math:`m~w.e.`].
+        new_ice_fraction (float): New volumetric ice fraction [-].
+    """
 
-    test_node.set_layer_temperature(new_temperature)
-    assert test_node.get_layer_temperature() == new_temperature
+    height = 0.1
+    density = 200.0
+    temperature = 270
+    lwc = 0.3
+    ice_fraction = 0.4
+    new_height = 1.0
+    new_temperature = 260.0
+    new_lwc = 0.5
+    new_ice_fraction = 0.9
+    new_refreeze = 0.01
 
-    test_node.set_layer_liquid_water_content(new_liquidwater_content)
-    assert test_node.get_layer_liquid_water_content() == new_liquidwater_content
+    def create_node(
+        self,
+        height: float = height,
+        density: float = density,
+        temperature: float = temperature,
+        lwc: float = lwc,
+        ice_fraction: float = ice_fraction,
+    ) -> Node:
+        """Instantiate a Node."""
 
-    test_node.set_layer_ice_fraction(new_ice_fraction)
-    assert test_node.get_layer_ice_fraction() == new_ice_fraction
+        node = Node(
+            height=height,
+            snow_density=density,
+            temperature=temperature,
+            liquid_water_content=lwc,
+            ice_fraction=ice_fraction,
+        )
+        return node
 
-    test_node.set_layer_refreeze(new_refreeze)
-    assert test_node.get_layer_refreeze() == new_refreeze
\ No newline at end of file
+    def test_create_node(self, conftest_boilerplate):
+        node = self.create_node()
+        assert isinstance(node, Node)
+        conftest_boilerplate.check_output(node.height, float, self.height)
+        conftest_boilerplate.check_output(
+            node.temperature, float, self.temperature
+        )
+        conftest_boilerplate.check_output(
+            node.liquid_water_content, float, self.lwc
+        )
+        conftest_boilerplate.check_output(
+            node.ice_fraction, float, self.ice_fraction
+        )
+        conftest_boilerplate.check_output(node.refreeze, float, 0.0)
+
+    def test_node_set_layer_height(self, conftest_boilerplate):
+        node = self.create_node()
+        node.set_layer_height(self.new_height)
+        conftest_boilerplate.check_output(
+            node.get_layer_height(), float, self.new_height
+        )
+
+    def test_node_set_layer_temperature(self, conftest_boilerplate):
+        node = self.create_node()
+        node.set_layer_temperature(self.new_temperature)
+        conftest_boilerplate.check_output(
+            node.get_layer_temperature(), float, self.new_temperature
+        )
+
+    def test_node_set_layer_liquid_water_content(self, conftest_boilerplate):
+        node = self.create_node()
+        node.set_layer_liquid_water_content(self.new_lwc)
+        conftest_boilerplate.check_output(
+            node.get_layer_liquid_water_content(), float, self.new_lwc
+        )
+
+    def test_node_set_layer_ice_fraction(self, conftest_boilerplate):
+        node = self.create_node()
+        node.set_layer_ice_fraction(self.new_ice_fraction)
+        conftest_boilerplate.check_output(
+            node.get_layer_ice_fraction(), float, self.new_ice_fraction
+        )
+
+    def test_node_set_layer_refreeze(self, conftest_boilerplate):
+        node = self.create_node()
+        node.set_layer_refreeze(self.new_refreeze)
+        conftest_boilerplate.check_output(
+            node.get_layer_refreeze(), float, self.new_refreeze
+        )
diff --git a/cosipy/tests/test_parameterisation_albedo.py b/cosipy/tests/test_parameterisation_albedo.py
index 556b322c..f94ef7ee 100644
--- a/cosipy/tests/test_parameterisation_albedo.py
+++ b/cosipy/tests/test_parameterisation_albedo.py
@@ -1,21 +1,100 @@
-from pytest import approx
-from cosipy.modules.albedo import updateAlbedo
-from cosipy.cpkernel.grid import *
+import numpy as np
+import pytest
+
+import constants
+import cosipy.modules.albedo as module_albedo
 from COSIPY import start_logging
 
-layer_heights = [0.1, 0.2, 0.3, 0.5, 0.5]
-layer_densities = [250, 250, 250, 917, 917]
-layer_temperatures = [260, 270, 271, 271.5, 272]
-layer_liquid_water = [0.0, 0.0, 0.0, 0.0, 0.0]
-#
-GRID = Grid(layer_heights, layer_densities, layer_temperatures, layer_liquid_water)
-GRID_ice = Grid(layer_heights[3:4], layer_densities[3:4], layer_temperatures[3:4],layer_liquid_water[3:4])
 
-def test_albedo_parameterisation():
+class TestParamAlbedoUpdate:
+    """Tests get/set methods for albedo properties."""
+
+    def test_updateAlbedo(self, conftest_mock_grid):
+        grid = conftest_mock_grid
+
+        surface_albedo, snow_albedo = module_albedo.updateAlbedo(
+            GRID=grid,
+            surface_temperature=270.0,
+            albedo_snow=constants.albedo_fresh_snow,
+        )
+        assert isinstance(surface_albedo, float)
+        assert constants.albedo_firn <= surface_albedo <= constants.albedo_fresh_snow
+        assert isinstance(snow_albedo, float)
+
+    def test_updateAlbedo_ice(self, conftest_mock_grid_ice, conftest_boilerplate):
+        grid_ice = conftest_mock_grid_ice
+        albedo, snow_albedo = module_albedo.updateAlbedo(
+            GRID=grid_ice,
+            surface_temperature=270.0,
+            albedo_snow=constants.albedo_fresh_snow,
+        )
+        assert conftest_boilerplate.check_output(albedo, float, constants.albedo_ice)
+        assert conftest_boilerplate.check_output(
+            snow_albedo, float, constants.albedo_fresh_snow
+        )
+
+
+class TestParamAlbedoSelection:
+    """Tests user selection of parametrisation method."""
+
+    @pytest.mark.parametrize("arg_method", ["Oerlemans98", "Bougamont05"])
+    def test_updateAlbedo_method(
+        self, monkeypatch, conftest_mock_grid, conftest_boilerplate, arg_method
+    ):
+        """Set method from constants.py when calculating albedo."""
+
+        grid = conftest_mock_grid
+
+        conftest_boilerplate.patch_variable(
+            monkeypatch, module_albedo.constants, {"albedo_method": arg_method}
+        )
+        assert constants.albedo_method == arg_method
+        surface_albedo, snow_albedo = module_albedo.updateAlbedo(
+            GRID=grid,
+            surface_temperature=270.0,
+            albedo_snow=constants.albedo_fresh_snow,
+        )
+        assert isinstance(surface_albedo, float)
+        assert isinstance(snow_albedo, float)
+
+    @pytest.mark.parametrize("arg_method", ["Wrong Method", "", None])
+    def test_updateAlbedo_method_error(
+        self, monkeypatch, conftest_mock_grid, conftest_boilerplate, arg_method
+    ):
+        grid = conftest_mock_grid
+        valid_methods = ["Oerlemans98", "Bougamont05"]
+
+        conftest_boilerplate.patch_variable(
+            monkeypatch, module_albedo.constants, {"albedo_method": arg_method}
+        )
+        assert constants.albedo_method == arg_method
+        error_message = (
+            f'Albedo method = "{constants.albedo_method}"',
+            f"is not allowed, must be one of",
+            f'{", ".join(valid_methods)}',
+        )
+
+        with pytest.raises(ValueError, match=" ".join(error_message)):
+            module_albedo.updateAlbedo(
+                GRID=grid,
+                surface_temperature=270.0,
+                albedo_snow=constants.albedo_fresh_snow,
+            )
+
+
+class TestParamAlbedoMethods:
+    """Tests methods for parametrising albedo."""
+
+    def test_method_Oerlemans(self, conftest_mock_grid):
+        """Get surface albedo without accounting for snow depth."""
 
-    albedo = updateAlbedo(GRID,3600)
-    assert albedo >= albedo_firn and albedo <= albedo_fresh_snow
+        grid = conftest_mock_grid
+        albedo_limit = constants.albedo_firn + (
+            constants.albedo_fresh_snow - constants.albedo_firn
+        ) * np.exp(0 / (constants.albedo_mod_snow_aging * 24.0))
 
-    albedo = updateAlbedo(GRID_ice,3600)
-    assert albedo == albedo_ice
+        compare_albedo = module_albedo.method_Oerlemans(grid)
 
+        assert isinstance(compare_albedo, float)
+        assert 0.0 <= compare_albedo <= 1.0
+        assert compare_albedo <= albedo_limit
diff --git a/cosipy/tests/test_parameterisation_densification.py b/cosipy/tests/test_parameterisation_densification.py
index 8f217b64..83402dc7 100644
--- a/cosipy/tests/test_parameterisation_densification.py
+++ b/cosipy/tests/test_parameterisation_densification.py
@@ -1,23 +1,51 @@
-from pytest import approx
-from cosipy.modules.densification import densification
-from cosipy.cpkernel.grid import *
+import numba
+import numpy as np
+
+# import constants
 from COSIPY import start_logging
+from cosipy.cpkernel.grid import Grid
+from cosipy.modules.densification import densification
+
+
+class TestParamDensification:
+    """Tests methods for parametrising densification.
+
+    Attributes:
+        heights (np.ndarray[float]): Snowpack heights for each layer
+            [:math:`m`].
+        densities (np.ndarray[float]): Snow densities for each layer
+            [:math:`kg~m^{-3}`].
+        temperatures (np.ndarray[float]): Temperatures for each layer
+            [:math:`K`].
+        liquid_water_contents (np.ndarray[float]): Liquid water content
+            for each layer [:math:`m~w.e.`].
+    """
+
+    # values are different to fixture
+    heights = numba.float64([10.0, 10.0, 10.0, 0.2, 0.2])
+    densities = numba.float64([450, 450, 450, 100, 100])
+    temperatures = numba.float64([260, 270, 271, 271.5, 272])
+    liquid_water_contents = numba.float64([0, 0, 0, 0, 0])
 
-layer_heights = [10.0, 10.0, 10.0, 0.2, 0.2]
-layer_densities = [450, 450, 450, 100, 100]
-layer_temperatures = [260, 270, 271, 271.5, 272]
-layer_liquid_water = [0.0, 0.0, 0.0, 0.0, 0.0]
+    def test_densification(self):
+        GRID = Grid(  # values are different to fixture
+            layer_heights=self.heights,
+            layer_densities=self.densities,
+            layer_temperatures=self.temperatures,
+            layer_liquid_water_content=self.liquid_water_contents,
+        )
+        SWE_before = np.array(GRID.get_height()) / np.array(GRID.get_density())
+        SWE_before_sum = np.nansum(SWE_before)
 
-GRID = Grid(layer_heights, layer_densities, layer_temperatures, layer_liquid_water)
-#
-def test_densification_parameterisation():
-    SWE_before = np.array(GRID.get_height()) / np.array(GRID.get_density())
-    SWE_before_sum = np.nansum(SWE_before)
+        densification(GRID=GRID, SLOPE=0.0, dt=3600)
 
-    densification(GRID, 0.0)
+        SWE_after = np.array(GRID.get_height()) / np.array(GRID.get_density())
+        SWE_after_sum = np.nansum(SWE_after)
 
-    SWE_after = np.array(GRID.get_height()) / np.array(GRID.get_density())
-    SWE_after_sum = np.nansum(SWE_after)
+        # Original assertions:
+        # assert np.isclose(SWE_before_sum, SWE_after_sum, atol=1e-3)
+        # assert SWE_before_sum == SWE_after_sum
 
-    #assert SWE_before_sum == approx(SWE_after_sum, abs=1e-3)
-    #assert SWE_before_sum == SWE_after_sum
+        # New assertion: Global densification method is set to "Boone",
+        # not "constant", so shouldn't the SWEs NOT match?
+        assert not np.isclose(SWE_before_sum, SWE_after_sum, atol=1e-3)
diff --git a/cosipy/tests/test_parameterisation_penetratingRadiation.py b/cosipy/tests/test_parameterisation_penetratingRadiation.py
index 2d3a85b7..9b11648a 100644
--- a/cosipy/tests/test_parameterisation_penetratingRadiation.py
+++ b/cosipy/tests/test_parameterisation_penetratingRadiation.py
@@ -1,1 +1,58 @@
-from pytest import approx
+import pytest
+import constants
+import cosipy.modules.penetratingRadiation as pRad
+
+
+class TestParamRadiation:
+    """Tests radiation methods."""
+
+    melt_water = 1.0
+    timedelta = 3600
+
+    def test_penetrating_radiation_error(
+        self,
+        monkeypatch,
+        conftest_mock_grid,
+        conftest_boilerplate,
+    ):
+        test_grid = conftest_mock_grid
+        allow_list = ["Bintanja95"]
+        conftest_boilerplate.patch_variable(
+            monkeypatch, pRad, {"penetrating_method": "WrongMethod"}
+        )
+        error_msg = (
+            f'Penetrating method = "{pRad.penetrating_method}" ',
+            f'is not allowed, must be one of {", ".join(allow_list)}',
+        )
+        with pytest.raises(ValueError, match="".join(error_msg)):
+            pRad.penetrating_radiation(GRID=test_grid, SWnet=800.0, dt=3600)
+
+    @pytest.mark.parametrize(
+        "arg_density", [250.0, constants.snow_ice_threshold + 1]
+    )
+    def test_method_Bintanja(
+        self,
+        monkeypatch,
+        conftest_mock_grid,
+        conftest_boilerplate,
+        arg_density,
+    ):
+        conftest_boilerplate.patch_variable(
+            monkeypatch, pRad, {"penetrating_method": "Bintanja95"}
+        )
+        test_grid = conftest_mock_grid
+        test_grid.add_fresh_snow(0.1, arg_density, 270.15, 0.0)
+        test_swnet = 800.0
+        if arg_density <= constants.snow_ice_threshold:
+            test_si = test_swnet * 0.1
+        else:
+            test_si = test_swnet * 0.2
+
+        melt_si = pRad.method_Bintanja(
+            GRID=test_grid, SWnet=test_swnet, dt=constants.dt
+        )
+        assert isinstance(melt_si, tuple)
+        compare_melt = melt_si[0]
+        assert isinstance(compare_melt, float)
+        compare_si = melt_si[1]
+        conftest_boilerplate.check_output(compare_si, float, test_si)
diff --git a/cosipy/tests/test_parameterisation_percolation.py b/cosipy/tests/test_parameterisation_percolation.py
index a409bcbc..2cc10967 100644
--- a/cosipy/tests/test_parameterisation_percolation.py
+++ b/cosipy/tests/test_parameterisation_percolation.py
@@ -1,17 +1,191 @@
-from pytest import approx
-from cosipy.cpkernel.grid import *
-from cosipy.modules.percolation import percolation
+import numpy as np
+import pytest
+
+# import constants
 from COSIPY import start_logging
+from cosipy.cpkernel.grid import Grid
+from cosipy.modules.percolation import percolation
+
+
+class TestParamPercolation:
+    """Tests percolation methods.
+
+    Attributes:
+        melt_water (float): Melt water at surface.
+        timedelta (int): Integration time, dt [s].
+    """
+
+    melt_water = 1.0
+    timedelta = 3600
+
+    def convert_mwe(self, mwe: float, depth: float) -> float:
+        """Converts m w.e. to kgm^-3."""
+        assert isinstance(mwe, float)
+        water = mwe / depth
+
+        return water
+
+    @pytest.mark.parametrize("arg_mwe", [0.0, 0.5, 1.0])
+    @pytest.mark.parametrize("arg_depth", [0.1, 1.0, 10.0])
+    def test_percolation_convert_mwe(self, arg_mwe: float, arg_depth: float):
+        compare_mwe = self.convert_mwe(mwe=arg_mwe, depth=arg_depth)
+        assert isinstance(compare_mwe, float)
+        if arg_depth < 1 and arg_mwe > 0:
+            assert compare_mwe > arg_mwe
+        else:
+            assert compare_mwe <= arg_mwe
+
+    def set_liquid_top_layer(self, grid: Grid, water: float):
+        """Set liquid water content for top layer."""
+        node_height_0 = grid.get_node_height(0)
+        water_in = self.convert_mwe(mwe=water, depth=node_height_0)
+        grid.set_node_liquid_water_content(
+            0, grid.get_node_liquid_water_content(0) + float(water_in)
+        )
+
+        return grid
+
+    @pytest.mark.parametrize("arg_water", [0.0, 0.5, 1.0])
+    @pytest.mark.parametrize("arg_depth", [0.1, 1.0, 10.0])
+    def test_set_liquid_top_layer(
+        self, conftest_mock_grid, arg_water: float, arg_depth: float
+    ):
+        GRID = conftest_mock_grid
+        water_in = self.convert_mwe(mwe=arg_water, depth=arg_depth)
+        assert isinstance(water_in, float)
+        start_water_layer_0 = GRID.get_node_liquid_water_content(0)
+        assert start_water_layer_0 == 0.0
+
+        GRID = self.set_liquid_top_layer(grid=GRID, water=water_in)
+
+        new_water_layer_0 = GRID.get_node_liquid_water_content(0)
+        if arg_water > 0:
+            assert new_water_layer_0 > 0.0
+            assert new_water_layer_0 > start_water_layer_0
+        else:
+            assert new_water_layer_0 == 0.0
+            assert new_water_layer_0 == start_water_layer_0
+
+    def get_layer_runoff(self, grid: Grid, idx: int) -> tuple:
+        """Calculate percolated runoff for a single layer."""
+        assert isinstance(idx, int)
+        theta_e = grid.get_node_irreducible_water_content(idx)
+        theta_w = grid.get_node_liquid_water_content(idx)
+        # Residual volume fraction of water (m^3 which is equal to m)
+        residual = np.maximum((theta_w - theta_e), 0.0)
+
+        return residual, theta_e, theta_w
+
+    @pytest.mark.parametrize("arg_idx", [0, 1, -1])
+    def test_get_layer_runoff(self, conftest_mock_grid, arg_idx: float):
+        GRID = conftest_mock_grid
+        theta_e = GRID.get_node_irreducible_water_content(arg_idx)
+        theta_w = GRID.get_node_liquid_water_content(arg_idx)
+
+        residual, new_theta_e, new_theta_w = self.get_layer_runoff(
+            grid=GRID, idx=int(arg_idx)
+        )
+
+        assert isinstance(residual, float)
+        assert np.isclose(residual, np.maximum((theta_w - theta_e), 0.0))
+        assert np.isclose(
+            GRID.get_node_irreducible_water_content(arg_idx), new_theta_e
+        )
+        assert np.isclose(
+            GRID.get_node_liquid_water_content(arg_idx), new_theta_w
+        )
+        assert np.isclose(
+            GRID.get_node_irreducible_water_content(arg_idx), theta_e
+        )
+        assert np.isclose(GRID.get_node_liquid_water_content(arg_idx), theta_w)
+
+    def percolate_layer(self, grid: Grid, idx: int) -> Grid:
+        residual, theta_e, theta_w = self.get_layer_runoff(grid=grid, idx=idx)
+        if residual > 0.0:
+            grid.set_node_liquid_water_content(idx, theta_w - residual)
+            residual = residual * grid.get_node_height(idx)
+            grid.set_node_liquid_water_content(
+                idx + 1,
+                grid.get_node_liquid_water_content(idx + 1)
+                + residual / grid.get_node_height(idx + 1),
+            )
+
+        return grid
+
+    @pytest.mark.parametrize("arg_idx", [0, 1, -2])
+    @pytest.mark.parametrize(
+        "arg_distribution",
+        ["zero", "random", "static", "decreasing", "increasing"],
+    )
+    def test_percolate_layer(
+        self,
+        conftest_mock_grid,
+        conftest_boilerplate,
+        arg_idx,
+        arg_distribution,
+    ):
+        """Tests percolation between two layers."""
+        GRID = conftest_mock_grid
+
+        GRID = conftest_boilerplate.regenerate_grid_values(
+            grid=GRID, distribution=arg_distribution
+        )
+        assert all(i >= 0.0 for i in GRID.get_liquid_water_content())
+        assert GRID.get_node_liquid_water_content(arg_idx) >= 0.0
+
+        start_theta_e = GRID.get_node_irreducible_water_content(arg_idx)
+        start_theta_w = GRID.get_node_liquid_water_content(arg_idx)
+        start_theta_e_next = GRID.get_node_irreducible_water_content(
+            arg_idx + 1
+        )
+        start_theta_w_next = GRID.get_node_liquid_water_content(arg_idx + 1)
+        start_residual = np.maximum((start_theta_w - start_theta_e), 0.0)
+        start_height = GRID.get_node_height(arg_idx)
+        height_next = GRID.get_node_height(arg_idx + 1)
+        predict_theta_w_next = start_theta_w_next + (
+            (start_residual * start_height) / height_next
+        )
+
+        GRID = self.percolate_layer(grid=GRID, idx=arg_idx)
+
+        assert all(i >= 0.0 for i in GRID.get_liquid_water_content())
+        assert GRID.get_node_liquid_water_content(arg_idx) >= 0.0
+
+        new_theta_e = GRID.get_node_irreducible_water_content(arg_idx)
+        new_theta_w = GRID.get_node_liquid_water_content(arg_idx)
+        new_theta_e_next = GRID.get_node_irreducible_water_content(arg_idx + 1)
+        new_theta_w_next = GRID.get_node_liquid_water_content(arg_idx + 1)
+        assert new_theta_w >= 0.0
+        assert np.isclose(new_theta_w, start_theta_w - start_residual)
+        assert np.isclose(new_theta_e, start_theta_e)
+        assert np.isclose(new_theta_e_next, start_theta_e_next)
+        assert np.isclose(new_theta_w_next, predict_theta_w_next)
+
+    @pytest.mark.parametrize("arg_melt", [0.0, 0.5, 1.0])
+    @pytest.mark.parametrize("arg_lwc", [0.0, 0.1, 0.5])
+    def test_percolation(self, conftest_mock_grid, arg_lwc, arg_melt):
+        GRID = conftest_mock_grid
+        heights = GRID.get_height()
+        lwc_array = np.full(GRID.number_nodes, arg_lwc)
+        GRID.set_liquid_water_content(lwc_array)
+        np.testing.assert_allclose(GRID.get_liquid_water_content(), lwc_array)
+
+        initial_mwe = (
+            np.nansum(np.array(GRID.get_liquid_water_content()) * heights)
+            + arg_melt
+        )
+
+        runoff = percolation(GRID, arg_melt, self.timedelta)
 
-layer_heights = [0.1, 0.2, 0.3, 0.5, 0.5]
-layer_densities = [250, 250, 250, 917, 917]
-layer_temperatures = [260, 270, 271, 271.5, 272]
-layer_liquid_water = [0.0, 0.0, 0.0, 0.0, 0.0]
+        final_mwe = (
+            np.nansum(np.array(GRID.get_liquid_water_content()) * heights)
+            + runoff
+        )
 
-GRID = Grid(layer_heights, layer_densities, layer_temperatures, layer_liquid_water)
-melt_water = 1.0
-timestamp = 7200
+        assert all(i >= 0.0 for i in GRID.get_liquid_water_content())
+        # all liquid water exits last node
+        last_node_lwc = GRID.grid[-1].liquid_water_content
+        assert last_node_lwc == 0.0
 
-def test_percolation_parameterisation():
-    q = percolation(GRID, melt_water, timestamp)
-    assert melt_water == np.nansum(GRID.get_liquid_water_content()) + q
\ No newline at end of file
+        assert runoff <= initial_mwe
+        assert np.isclose(initial_mwe, final_mwe)
diff --git a/cosipy/tests/test_parameterisation_refreezing.py b/cosipy/tests/test_parameterisation_refreezing.py
index 5cf6d0ea..a7b43965 100644
--- a/cosipy/tests/test_parameterisation_refreezing.py
+++ b/cosipy/tests/test_parameterisation_refreezing.py
@@ -1,16 +1,43 @@
-from pytest import approx
-from cosipy.cpkernel.grid import *
-from cosipy.modules.refreezing import refreezing
+import numba
+import numpy as np
+
+# import constants
 from COSIPY import start_logging
+from cosipy.cpkernel.grid import Grid
+from cosipy.modules.refreezing import refreezing
+
+
+class TestParamRefreezing:
+    """Tests methods for parametrising refreezing.
 
-layer_heights = [0.1, 0.2, 0.3, 0.5, 0.5]
-layer_densities = [250, 250, 250, 917, 917]
-layer_temperatures = [260, 270, 271, 271.5, 272]
-layer_liquid_water = [0.01, 0.01, 0.01, 0.01, 0.01]
+    Attributes:
+        heights (np.ndarray[float]): Snowpack heights for each layer
+            [:math:`m`].
+        densities (np.ndarray[float]): Snow densities for each layer
+            [:math:`kg~m^{-3}`].
+        temperatures (np.ndarray[float]): Temperatures for each layer
+            [:math:`K`].
+        liquid_water_contents (np.ndarray[float]): Liquid water content
+            for each layer [:math:`m~w.e.`].
+    """
 
-GRID = Grid(layer_heights, layer_densities, layer_temperatures, layer_liquid_water)
+    # values are different to fixture
+    heights = numba.float64([0.1, 0.2, 0.3, 0.5, 0.5])
+    densities = numba.float64([250, 250, 250, 917, 917])
+    temperatures = numba.float64([260, 270, 271, 271.5, 272])
+    liquid_water_contents = numba.float64([0.01, 0.01, 0.01, 0.01, 0.01])
 
-def test_refreezing_parametrisation():
-    water_content_before = np.nansum(GRID.get_liquid_water_content())
-    water_refreezed = refreezing(GRID)
-    #assert water_content_before == np.nansum(GRID.get_liquid_water_content()) + water_refreezed
\ No newline at end of file
+    def test_refreezing(self):
+        GRID = Grid(
+            layer_heights=self.heights,
+            layer_densities=self.densities,
+            layer_temperatures=self.temperatures,
+            layer_liquid_water_content=self.liquid_water_contents,
+        )
+        water_content_before = np.nansum(GRID.get_liquid_water_content())
+        water_refreezed = refreezing(GRID)
+        water_content_after = (
+            np.nansum(GRID.get_liquid_water_content()) + water_refreezed
+        )
+        # Bug?
+        assert np.isclose(water_content_before, water_content_after)
diff --git a/cosipy/tests/test_parameterisation_roughness.py b/cosipy/tests/test_parameterisation_roughness.py
index a176ecb6..e6c1c4a8 100644
--- a/cosipy/tests/test_parameterisation_roughness.py
+++ b/cosipy/tests/test_parameterisation_roughness.py
@@ -1,20 +1,22 @@
-from pytest import approx
-from cosipy.modules.roughness import updateRoughness
-from cosipy.cpkernel.grid import *
+import numpy as np
+
+import constants
 from COSIPY import start_logging
+from cosipy.modules.roughness import updateRoughness
 
-layer_heights = [0.1, 0.2, 0.3, 0.5, 0.5]
-layer_densities = [250, 250, 250, 917, 917]
-layer_temperatures = [260, 270, 271, 271.5, 272]
-layer_liquid_water = [0.0, 0.0, 0.0, 0.0, 0.0]
-#
-GRID = Grid(layer_heights, layer_densities, layer_temperatures, layer_liquid_water)
-GRID_ice = Grid(layer_heights[3:4], layer_densities[3:4], layer_temperatures[3:4],layer_liquid_water[3:4])
 
-def test_roughness_parameterisation():
+class TestParamRoughness:
+    """Tests methods for parametrising roughness."""
 
-    roughness = updateRoughness(GRID,3600)
-    assert roughness >= roughness_fresh_snow/1000 and roughness <= roughness_firn/1000
+    def test_updateRoughness(self, conftest_mock_grid, conftest_mock_grid_ice):
+        GRID = conftest_mock_grid
+        roughness = updateRoughness(GRID)
+        assert (
+            constants.roughness_fresh_snow / 1000
+            <= roughness
+            <= constants.roughness_firn / 1000
+        )
 
-    roughness = updateRoughness(GRID_ice,3600)
-    assert roughness == roughness_ice/1000
\ No newline at end of file
+        GRID_ice = conftest_mock_grid_ice
+        ice_roughness = updateRoughness(GRID_ice)
+        assert np.isclose(ice_roughness, constants.roughness_ice / 1000)
diff --git a/cosipy/tests/test_parameterisation_surfaceTemperature.py b/cosipy/tests/test_parameterisation_surfaceTemperature.py
index bbfa9a10..83846ae4 100644
--- a/cosipy/tests/test_parameterisation_surfaceTemperature.py
+++ b/cosipy/tests/test_parameterisation_surfaceTemperature.py
@@ -1,26 +1,50 @@
-from pytest import approx
-from cosipy.cpkernel.grid import *
-from cosipy.modules.surfaceTemperature import update_surface_temperature
+import constants
 from COSIPY import start_logging
+from cosipy.modules.surfaceTemperature import update_surface_temperature
 
 
-layer_heights = [0.1, 0.2, 0.3, 0.5, 0.5]
-layer_densities = [250, 250, 250, 917, 917]
-layer_temperatures = [260, 270, 271, 271.5, 272]
-layer_liquid_water = [0.0, 0.0, 0.0, 0.0, 0.0]
-
-GRID = Grid(layer_heights, layer_densities, layer_temperatures, layer_liquid_water)
+class TestParamSurfaceTemperature:
+    """Tests methods for parametrising surface temperature."""
 
-def test_surface_Temperature_parameterisation():
+    def test_surface_Temperature_parameterisation(self, conftest_mock_grid):
+        GRID = conftest_mock_grid
 
-    fun, surface_temperature, lw_radiation_in, lw_radiation_out, sensible_heat_flux, latent_heat_flux, \
-    ground_heat_flux, sw_radiation_net, rho, Lv, Cs_t, Cs_q, q0, q2 \
-        = update_surface_temperature(GRID, 0.6, (0.24/1000), 275, 0.6, 789, 1000, 4.5, 0.0, 0.1)
+        (
+            fun,
+            surface_temperature,
+            lw_radiation_in,
+            lw_radiation_out,
+            sensible_heat_flux,
+            latent_heat_flux,
+            ground_heat_flux,
+            rain_heat_flux,
+            rho,
+            Lv,
+            monin_obukhov_length,
+            Cs_t,
+            Cs_q,
+            q0,
+            q2,
+        ) = update_surface_temperature(
+            # Old args: GRID, 0.6, (0.24 / 1000), 275, 0.6, 789, 1000, 4.5, 0.0, 0.1
+            GRID=GRID,
+            dt=3600,
+            z=2,
+            z0=(0.24 / 1000),
+            T2=275,
+            rH2=50,
+            p=1000,
+            SWnet=789,
+            u2=3.5,
+            RAIN=0.1,
+            SLOPE=0.0,
+            LWin=None,
+            N=0.5,  # otherwise tries to retrieve LWin from non-existent file
+        )
 
-    assert surface_temperature <= zero_temperature and surface_temperature >= 220.0
-    assert lw_radiation_in <= 400 and lw_radiation_in >= 0
-    assert lw_radiation_out <= 0 and lw_radiation_out >= -400
-    assert sensible_heat_flux <= 250 and sensible_heat_flux >= -250
-    assert latent_heat_flux <= 200 and latent_heat_flux >= -200
-    assert ground_heat_flux <= 100 and ground_heat_flux >= -100
-    assert sw_radiation_net <= 1000 and sw_radiation_net >= 0
+        assert constants.zero_temperature >= surface_temperature >= 220.0
+        assert 400 >= lw_radiation_in >= 0
+        assert 0 >= lw_radiation_out >= -400
+        assert 250 >= sensible_heat_flux >= -250
+        assert 200 >= latent_heat_flux >= -200
+        assert 100 >= ground_heat_flux >= -100
diff --git a/cosipy/tests/test_postprocessing_field_plots.py b/cosipy/tests/test_postprocessing_field_plots.py
new file mode 100644
index 00000000..8bffe5e9
--- /dev/null
+++ b/cosipy/tests/test_postprocessing_field_plots.py
@@ -0,0 +1,368 @@
+"""Tests for postprocessing of field plots.
+
+Call ``plt.close("all")`` at the end of any test that creates a
+matplotlib object.
+"""
+
+import argparse
+from unittest.mock import patch
+
+import cartopy
+import cartopy.crs as ccrs
+import matplotlib
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+import pytest
+import xarray as xr
+
+import postprocessing.field_plots.plot_cosipy_fields as pcf
+
+# import constants
+# from COSIPY import start_logging
+
+
+class TestPostprocessPlotFieldsHandling:
+    """Tests data handling before passing to plots."""
+
+    @pytest.mark.parametrize("arg_mean", [True, False])
+    def test_get_selection(self, conftest_mock_xr_dataset, arg_mean):
+        dataset = conftest_mock_xr_dataset.copy()
+        if arg_mean:
+            timestamp = "2009-01-01"  # see test_get_selection_error
+        else:
+            timestamp = "2009-01-01T12:00:00"
+
+        compare_array = pcf.get_selection(
+            array=dataset["TS"], timestamp=timestamp, mean=arg_mean
+        )
+
+        assert isinstance(compare_array, xr.DataArray)
+        if not arg_mean:
+            assert compare_array.time == dataset.time[0]
+        else:
+            assert "time" not in compare_array.dims
+
+    def test_get_selection_error(self, conftest_mock_xr_dataset):
+        dataset = conftest_mock_xr_dataset.copy()
+        timestamp = "2009-01-01"
+        error_message = (
+            "Selected time index refers to an entire day,",
+            "not a single timestep.",
+            "Use `--mean` to plot the daily mean.",
+        )
+        with pytest.raises(KeyError, match=" ".join(error_message)):
+            pcf.get_selection(
+                array=dataset["TS"], timestamp=timestamp, mean=False
+            )
+
+    def test_check_2d(self):
+        dims = {}
+        reference_time = pd.Timestamp("2009-01-01T12:00:00")
+        dims["name"] = ["time", "lat", "lon"]
+        dims["time"] = pd.date_range(reference_time, periods=4, freq="6H")
+        elevation = xr.Variable(
+            data=1000 + 10 * np.random.rand(2, 1),
+            dims=dims["name"][1:],
+            attrs={"long_name": "Elevation", "units": "m"},
+        )
+        dataset = xr.Dataset(
+            data_vars=dict(HGT=elevation),
+            coords=dict(
+                time=dims["time"],
+                lat=(["lat"], [30.460, 30.463]),
+                lon=(["lon"], [90.621]),
+                reference_time=dims["time"][0],
+            ),
+            attrs=dict(
+                description="Weather related data.",
+                Full_fiels="True",  # match typo in io.py
+            ),
+        )
+        error_message = "Spatial coordinates are not 2D."
+        with pytest.raises(ValueError, match=error_message):
+            pcf.check_2d(dataset)
+
+
+class TestPostprocessPlotFieldsPlotting:
+    """Tests data plotting."""
+
+    @pytest.mark.parametrize("arg_x", [1.00, 1.54, 1.21, 1.01])
+    @pytest.mark.parametrize("arg_tex", [True, False])
+    def test_set_contour_format(self, arg_x: float, arg_tex: bool):
+        test_label = f"{arg_x:.1f}"
+        if test_label.endswith("0"):
+            test_label = f"{arg_x:.0f}"
+
+        matplotlib.rcParams["text.usetex"] = arg_tex
+        if not arg_tex:
+            assert not matplotlib.rcParams["text.usetex"]
+        else:
+            assert matplotlib.rcParams["text.usetex"]
+
+        compare_label = pcf.set_contour_format(x=arg_x)
+        assert f"{test_label} m" == compare_label
+
+        matplotlib.rcParams.update(matplotlib.rcParamsDefault)
+        plt.close("all")
+
+    def test_set_gridlines(self):
+        ax = plt.subplot(projection=ccrs.PlateCarree())
+
+        compare_ax = pcf.set_gridlines(ax=ax)
+        assert isinstance(compare_ax, cartopy.mpl.geoaxes.GeoAxes)
+        assert compare_ax._gridliners
+        for gridlines in compare_ax._gridliners:
+            kwargs = gridlines.collection_kwargs
+            assert kwargs.get("linewidth", None) == 0.5
+            assert kwargs.get("color", None) == "gray"
+            assert kwargs.get("alpha", None) == 0.1
+            assert kwargs.get("linestyle", None) == "--"
+
+        plt.close("all")
+
+    def test_plot_topography(self, conftest_mock_xr_dataset):
+        dataset = conftest_mock_xr_dataset.copy()
+        ax = plt.subplot(projection=ccrs.PlateCarree())
+
+        compare_ax = pcf.plot_topography(ax=ax, elevation=dataset.HGT)
+        assert isinstance(compare_ax, plt.Axes)
+
+        compare_interval = ax.xaxis.properties()["data_interval"]
+        assert np.isclose(compare_interval[0], dataset.lon.min())
+        assert np.isclose(compare_interval[1], dataset.lon.max())
+
+        compare_interval = ax.yaxis.properties()["data_interval"]
+        assert np.isclose(compare_interval[0], dataset.lat.min())
+        assert np.isclose(compare_interval[1], dataset.lat.max())
+
+        plt.close("all")
+
+    @pytest.mark.parametrize("arg_name", ["TS", "HGT"])
+    @pytest.mark.parametrize("arg_topography", [True, False])
+    @pytest.mark.parametrize("arg_plot_type", ["contour", "mesh"])
+    @pytest.mark.parametrize("arg_mean", [True, False])
+    def test_plot_axes(
+        self,
+        conftest_mock_xr_dataset,
+        conftest_boilerplate,
+        arg_name: str,
+        arg_topography: bool,
+        arg_plot_type: str,
+        arg_mean: bool,
+    ):
+        dataset = conftest_mock_xr_dataset.copy()
+        test_array = dataset[arg_name]
+        test_ax = plt.subplot(projection=ccrs.PlateCarree())
+        if arg_topography:
+            topography = dataset.HGT
+        else:
+            topography = None
+        timestamp = conftest_boilerplate.set_timestamp(day=arg_mean)
+
+        compare_ax = pcf.plot_axes(
+            ax=test_ax,
+            array=test_array,
+            timestamp=timestamp,
+            topography=topography,
+            plot_type=arg_plot_type,
+            mean=arg_mean,
+        )
+
+        assert isinstance(compare_ax, plt.Axes)
+        assert compare_ax._gridliners
+        assert compare_ax.get_title() == dataset[arg_name].long_name.title()
+
+        plt.close("all")
+
+    def test_plot_axes_error(self, conftest_mock_xr_dataset):
+        dataset = conftest_mock_xr_dataset.copy()
+        test_ax = plt.subplot(projection=ccrs.PlateCarree())
+        test_plot_type = "missing_type"
+
+        error_message = (
+            f"{test_plot_type} is an invalid plot type.",
+            f"Use 'mesh' or 'contour'.",
+        )
+        with pytest.raises(ValueError, match=" ".join(error_message)):
+            pcf.plot_axes(
+                ax=test_ax,
+                array=dataset.TS,
+                timestamp="2009-01-01T12:00:00",
+                topography=None,
+                plot_type=test_plot_type,
+                mean=False,
+            )
+
+        plt.close("all")
+
+    @pytest.mark.parametrize("arg_list", [[], None])
+    def test_create_subplots_for_variables_error(self, arg_list):
+        error_message = "Passed empty variable array to adaptive subplots."
+        with pytest.raises(ValueError, match=error_message):
+            pcf.create_subplots_for_variables(variables=arg_list)
+
+        plt.close("all")
+
+    def test_create_subplots_for_variables(self):
+        test_variables = [["a", "b", "c"], ["d", "e", "f"]]
+
+        compare_figure, compare_axes = pcf.create_subplots_for_variables(
+            variables=test_variables
+        )
+        assert isinstance(compare_figure, plt.Figure)
+        assert isinstance(compare_axes, np.ndarray)
+        assert compare_axes.shape == (2, 3)  # match test_variables
+        for rows in compare_axes:
+            assert isinstance(rows, np.ndarray)
+            assert rows.shape == (len(test_variables[0]),)
+            for subplot in rows:
+                assert isinstance(subplot, cartopy.mpl.geoaxes.GeoAxesSubplot)
+                assert subplot.projection == ccrs.PlateCarree()
+
+        plt.close("all")
+
+    @pytest.mark.parametrize("arg_mean", [True, False])
+    @pytest.mark.parametrize("arg_plot_type", ["contour", "mesh"])
+    def test_plot_data_short_name(
+        self,
+        conftest_mock_xr_dataset,
+        conftest_boilerplate,
+        arg_mean,
+        arg_plot_type,
+    ):
+        dataset = conftest_mock_xr_dataset.copy(deep=True)
+        timestamp = conftest_boilerplate.set_timestamp(day=arg_mean)
+        compare_figure, compare_axes = pcf.plot_data(
+            data=dataset,
+            timestamp=timestamp,
+            short_name="TS",
+            mean=arg_mean,
+            variables=None,
+            plot_type=arg_plot_type,
+        )
+        compare_params = {
+            "plot": (compare_figure, compare_axes),
+            "title": dataset["TS"].long_name.title(),
+        }
+
+        conftest_boilerplate.check_plot(plot_params=compare_params)
+
+        plt.close("all")
+
+    @pytest.mark.parametrize("arg_mean", [True, False])
+    @pytest.mark.parametrize("arg_plot_type", ["contour", "mesh"])
+    def test_plot_data_variables(
+        self,
+        conftest_mock_xr_dataset,
+        conftest_boilerplate,
+        arg_mean,
+        arg_plot_type,
+    ):
+        dataset = conftest_mock_xr_dataset.copy(deep=True)
+        timestamp = conftest_boilerplate.set_timestamp(day=arg_mean)
+        test_variables = [
+            ["TS", "HGT"],
+            ["TS", "HGT"],
+        ]
+
+        compare_figure, compare_axes = pcf.plot_data(
+            data=dataset,
+            timestamp=timestamp,
+            short_name=None,
+            mean=arg_mean,
+            variables=test_variables,
+            plot_type=arg_plot_type,
+        )
+        assert isinstance(compare_figure, plt.Figure)
+        assert isinstance(compare_axes, np.ndarray)
+        assert compare_axes.shape == (2, 2)
+        compare_variables = []
+        for rows in compare_axes:
+            assert isinstance(rows, np.ndarray)
+            for subplot in rows:
+                assert isinstance(subplot, cartopy.mpl.geoaxes.GeoAxes)
+                compare_variables.append(subplot.get_title())
+
+        for title in compare_variables:
+            assert isinstance(title, str)
+            assert title in ["Surface Temperature", "Elevation"]
+
+        plt.close("all")
+
+    @pytest.mark.parametrize("arg_mean", [True, False])
+    def test_plotMesh(
+        self,
+        conftest_mock_open_dataset,
+        conftest_hide_plot,
+        conftest_boilerplate,
+        arg_mean,
+    ):
+        timestamp = conftest_boilerplate.set_timestamp(day=arg_mean)
+        _ = conftest_mock_open_dataset  # use mock data instead of file
+        _ = conftest_hide_plot  # avoid user interaction
+
+        pcf.plotMesh(
+            filename="./path/file",
+            pdate=timestamp,
+            var="TS",
+            mean=arg_mean,
+            save=False,
+        )
+
+        compare_figure = plt.gcf()
+        compare_axes = plt.gca()
+        compare_params = {
+            "plot": (compare_figure, compare_axes),
+            "title": "Surface Temperature",
+        }
+        conftest_boilerplate.check_plot(plot_params=compare_params)
+
+        plt.close("all")
+
+    @pytest.mark.parametrize("arg_mean", [True, False])
+    def test_plotContour(
+        self,
+        conftest_mock_open_dataset,
+        conftest_hide_plot,
+        conftest_boilerplate,
+        arg_mean,
+    ):
+        timestamp = conftest_boilerplate.set_timestamp(day=arg_mean)
+        _ = conftest_mock_open_dataset  # use mock data instead of file
+        _ = conftest_hide_plot  # avoid user interaction
+
+        pcf.plotContour(
+            filename="./path/file",
+            pdate=timestamp,
+            var="TS",
+            mean=arg_mean,
+            save=False,
+        )
+
+        compare_figure = plt.gcf()
+        compare_axes = plt.gca()
+        compare_params = {
+            "plot": (compare_figure, compare_axes),
+            "title": "Surface Temperature",
+        }
+        conftest_boilerplate.check_plot(plot_params=compare_params)
+
+        plt.close("all")
+
+    def test_parse_arguments(self):
+        with patch(
+            "sys.argv",
+            ["main", "--file", "./path/file", "--date", "'2009-01-01'"],
+        ):
+            args = pcf.parse_arguments()
+        assert isinstance(args, argparse.Namespace)
+
+        # Required
+        assert args.file == "./path/file"
+        assert args.pdate == "'2009-01-01'"
+        # Defaults
+        assert args.variable is None
+        assert isinstance(args.plot_type, int) and args.plot_type == 1
+        assert isinstance(args.mean, bool) and not args.mean
+        assert isinstance(args.save, bool) and not args.save