improve readibility of the DiffusionTensor class and add exceptions

features/diffusion_tensor.feature

@@ -15,3 +15,4 @@ Scenario: Generate the diffusion tensor
     And the expected diffusion tensor at a given time is computed
     And the update fails if the time is not between the lowest and largest times contained in the time vector
     And the update at a given time of the steady diffusion tensor does not change the values
+    And the initialization fails if the inputs do not have the proper type

features/diffusion_term.feature

@@ -11,3 +11,4 @@ Scenario: Generate the diffusion term of the convection-diffusion PDE as a linea
     And the result of the matrice vector product of the linear operator has the expected values
     And the linear operator is updated as expected at a given time
     And the update at a given time of the steady linear operator does not change the values
+    And the initialization fails if the diffusion tensor is not an object of the class DiffusionTensor

src/pheromone_dispersion/diffusion_operator.py

 import numpy as np
 from scipy.sparse.linalg import LinearOperator as LinOp
 
+from pheromone_dispersion.diffusion_tensor import _validate_diffusion_tensor
+
 """
 Module that contains the implementation of the linear operator of the diffusion term of the pheromone propagation model.
 """
@@ -40,6 +42,7 @@ class Diffusion(LinOp):
         msh : ~pheromone_dispersion.geom.MeshRect2D
             The geometry of the domain.
         """
+        _validate_diffusion_tensor(K, msh)
         self.msh = msh
         self.K = K
         self.shape = (self.msh.y.size * self.msh.x.size, self.msh.y.size * self.msh.x.size)

src/pheromone_dispersion/diffusion_tensor.py

 import numpy as np
 
+from pheromone_dispersion.velocity import Velocity
+from pheromone_dispersion.velocity import _validate_velocity_field
+
 """
 Module that contains the implementation of the diffusion term diffusion tensor.
 """
@@ -14,7 +17,7 @@ class DiffusionTensor:
     a diffusion coefficient :math:`K_{u^\perp}` in the crosswind direction.
     Hence, the anistropic diffusion tensor is given by
     :math:`\mathbf{K} = R(\vec{u})diag(K_u,K_{u^\perp})R(\vec{u})^T`
-    with :math:`R(\vec{u})` the rotation matrice
+    with :math:`R(\vec{u})` the rotation matrix
     of angle :math:`\theta` between the wind field and the cartesian frame.
 
     Attributes
@@ -46,6 +49,7 @@ class DiffusionTensor:
         """
         # To do
         # add exceptions to make sure that all the inputs have the right type
+        _validate_inputs(U, K_u, K_u_t)
         self.U = U
         self.K_u = K_u
         self.K_u_t = K_u_t
@@ -60,42 +64,24 @@ class DiffusionTensor:
         """
 
         # On the vertical interfaces of the mesh
-        norm_U_at_vertical_interface = np.linalg.norm(self.U.at_vertical_interface, axis=2)
-        U_at_vertical_interface = self.U.at_vertical_interface / norm_U_at_vertical_interface[:, :, None]
-        self.at_vertical_interface = np.zeros((U_at_vertical_interface.shape[0], U_at_vertical_interface.shape[1], 2, 2))
-        self.at_vertical_interface[:, :, 0, 0] = self.K_u_t
-        self.at_vertical_interface[:, :, 1, 1] = self.K_u_t
-        self.at_vertical_interface[:, :, 0, 0] += (
-            -(self.K_u_t - self.K_u) * U_at_vertical_interface[:, :, 0] * U_at_vertical_interface[:, :, 0]
-        )
-        self.at_vertical_interface[:, :, 0, 1] += (
-            -(self.K_u_t - self.K_u) * U_at_vertical_interface[:, :, 0] * U_at_vertical_interface[:, :, 1]
-        )
-        self.at_vertical_interface[:, :, 1, 0] += (
-            -(self.K_u_t - self.K_u) * U_at_vertical_interface[:, :, 0] * U_at_vertical_interface[:, :, 1]
-        )
-        self.at_vertical_interface[:, :, 1, 1] += (
-            -(self.K_u_t - self.K_u) * U_at_vertical_interface[:, :, 1] * U_at_vertical_interface[:, :, 1]
-        )
+        self.at_vertical_interface = self._compute_diffusion_tensor_from_velocity_field_at_interface(self.U.at_vertical_interface)
 
         # On the horizontal interfaces of the mesh
-        norm_U_at_horizontal_interface = np.linalg.norm(self.U.at_horizontal_interface, axis=2)
-        U_at_horizontal_interface = self.U.at_horizontal_interface / norm_U_at_horizontal_interface[:, :, None]
-        self.at_horizontal_interface = np.zeros((U_at_horizontal_interface.shape[0], U_at_horizontal_interface.shape[1], 2, 2))
-        self.at_horizontal_interface[:, :, 0, 0] = self.K_u_t
-        self.at_horizontal_interface[:, :, 1, 1] = self.K_u_t
-        self.at_horizontal_interface[:, :, 0, 0] += (
-            -(self.K_u_t - self.K_u) * U_at_horizontal_interface[:, :, 0] * U_at_horizontal_interface[:, :, 0]
-        )
-        self.at_horizontal_interface[:, :, 0, 1] += (
-            -(self.K_u_t - self.K_u) * U_at_horizontal_interface[:, :, 0] * U_at_horizontal_interface[:, :, 1]
-        )
-        self.at_horizontal_interface[:, :, 1, 0] += (
-            -(self.K_u_t - self.K_u) * U_at_horizontal_interface[:, :, 0] * U_at_horizontal_interface[:, :, 1]
-        )
-        self.at_horizontal_interface[:, :, 1, 1] += (
-            -(self.K_u_t - self.K_u) * U_at_horizontal_interface[:, :, 1] * U_at_horizontal_interface[:, :, 1]
-        )
+        self.at_horizontal_interface = self._compute_diffusion_tensor_from_velocity_field_at_interface(self.U.at_horizontal_interface)
+
+    def _compute_diffusion_tensor_from_velocity_field_at_interface(self, U_i):
+        # Compute the diffusion tensor at an interface given the velocity field at this interface
+        norm_U_i = np.linalg.norm(U_i, axis=2, keepdims=True)
+        U_i_normalized = np.divide(U_i, norm_U_i, out=np.zeros_like(U_i), where=norm_U_i != 0)
+        K_tensor_i = np.zeros(U_i.shape[:2] + (2, 2))
+        K_tensor_i[:, :, 0, 0] = self.K_u_t
+        K_tensor_i[:, :, 1, 1] = self.K_u_t
+        delta_K = self.K_u - self.K_u_t
+        K_tensor_i[:, :, 0, 0] += delta_K * U_i_normalized[:, :, 0] * U_i_normalized[:, :, 0]
+        K_tensor_i[:, :, 0, 1] += delta_K * U_i_normalized[:, :, 0] * U_i_normalized[:, :, 1]
+        K_tensor_i[:, :, 1, 0] += delta_K * U_i_normalized[:, :, 1] * U_i_normalized[:, :, 0]
+        K_tensor_i[:, :, 1, 1] += delta_K * U_i_normalized[:, :, 1] * U_i_normalized[:, :, 1]
+        return K_tensor_i
 
     def at_current_time(self, tc):
         """
@@ -116,6 +102,26 @@ class DiffusionTensor:
         and :attr:`at_horizontal_interface` from the updated velocity field
         using the method :meth:`diffusion_tensor_from_velocity_field`.
         """
+        if self.U.t is not None:
+            self.U.at_current_time(tc)
+            self.diffusion_tensor_from_velocity_field()
 
-        self.U.at_current_time(tc)
-        self.diffusion_tensor_from_velocity_field()
+
+def _validate_inputs(U, K_u, K_u_t):
+    # Validate the inputs of the DiffusionTensor class and ensure correct types.
+    if not isinstance(U, Velocity):
+        raise TypeError('The provided velocity field is not of type Velocity.')
+    if isinstance(K_u, np.ndarray) and K_u.size == 1:
+        K_u = K_u.item()
+    if isinstance(K_u_t, np.ndarray) and K_u_t.size == 1:
+        K_u_t = K_u_t.item()
+    if not isinstance(K_u, (int, float)) or not isinstance(K_u_t, (int, float)):
+        raise TypeError('The provided diffusion coefficients are neither float or int.')
+
+
+def _validate_diffusion_tensor(K, msh):
+    # Check that the diffusion tensor is of class DiffusionTensor
+    # and that the shape of the attributes match with the mesh.
+    if not isinstance(K, DiffusionTensor):
+        raise TypeError('The provided diffusion tensor is not of type DiffusionTensor.')
+    _validate_velocity_field(K.U, msh)

tests/test_diffusion_tensor.py

@@ -78,6 +78,8 @@ def the_diffusion_tensor_at_the_vertical_interfaces_has_the_expected_shape(K, U)
 @then("the expected diffusion tensor at the vertical interfaces is computed")
 def the_expected_diffusion_tensor_at_the_vertical_interfaces_is_computed(K):
     """the expected diffusion tensor at the vertical interfaces is computed"""
+    # print(K.at_vertical_interface)
+    # print(K.K_u, K.K_u_t)
     truth = np.ones((3, 3, 2, 2))
     truth[:, :, 0, 0] = 5.5
     truth[:, :, 1, 1] = 5.5
@@ -137,3 +139,17 @@ def the_update_at_a_given_time_of_the_steady_diffusion_tensor_does_not_change_th
     K_steady.at_current_time(1.5)
     assert (K_steady_Hori_save == K_steady.at_horizontal_interface).all()
     assert (K_steady_Vert_save == K_steady.at_vertical_interface).all()
+
+
+@then("the initialization fails if the inputs do not have the proper type")
+def the_initialization_fails_if_the_inputs_do_not_have_the_proper_type(U, K_u, K_u_t):
+    """the initialization fails if the inputs do not have the proper type"""
+    with pytest.raises(TypeError) as e:
+        DiffusionTensor(U, [False, False], K_u_t)
+    assert e.type == TypeError
+    with pytest.raises(TypeError) as e:
+        DiffusionTensor(U, K_u, [False, False])
+    assert e.type == TypeError
+    with pytest.raises(TypeError) as e:
+        DiffusionTensor(3.0, K_u, K_u_t)
+    assert e.type == TypeError

tests/test_diffusion_term.py

 from pathlib import Path
 
 import numpy as np
+import pytest
 from pytest_bdd import given
 from pytest_bdd import scenario
 from pytest_bdd import then
@@ -113,3 +114,11 @@ def the_update_at_a_given_time_of_the_steady_linear_operator_does_not_change_the
     diffusion_steady.at_current_time(1.5)
     assert (diffusion_steady_K_hori == diffusion_steady.K.at_horizontal_interface).all()
     assert (diffusion_steady_K_vert == diffusion_steady.K.at_vertical_interface).all()
+
+
+@then("the initialization fails if the diffusion tensor is not an object of the class DiffusionTensor")
+def the_initialization_fails_if_the_diffusion_tensor_is_not_an_object_of_the_class_DiffusionTensor(msh):
+    """the initialization fails if the diffusion tensor is not an object of the class DiffusionTensor"""
+    with pytest.raises(TypeError) as e:
+        Diffusion(3.0, msh)
+    assert e.type == TypeError