Move calculating current to FDTDGrid member funcion

This calculation only uses information from the grid so makes sense to
be done by the grid class. The previous helper function was passed a
copy of the grid as a parameter already.

gprMax/fields_outputs.py

@@ -140,54 +140,3 @@ def write_hd5_data(basegrp, grid, is_subgrid=False):
 
         for output in rx.outputs:
             basegrp["rxs/rx" + str(rxindex + 1) + "/" + output] = rx.outputs[output]
-
-
-def Ix(x, y, z, Hx, Hy, Hz, G):
-    """Calculates the x-component of current at a grid position.
-
-    Args:
-        x, y, z: floats for coordinates of position in grid.
-        Hx, Hy, Hz: numpy array of magnetic field values.
-        G: FDTDGrid class describing a grid in a model.
-    """
-
-    if y == 0 or z == 0:
-        Ix = 0
-    else:
-        Ix = G.dy * (Hy[x, y, z - 1] - Hy[x, y, z]) + G.dz * (Hz[x, y, z] - Hz[x, y - 1, z])
-
-    return Ix
-
-
-def Iy(x, y, z, Hx, Hy, Hz, G):
-    """Calculates the y-component of current at a grid position.
-
-    Args:
-        x, y, z: floats for coordinates of position in grid.
-        Hx, Hy, Hz: numpy array of magnetic field values.
-        G: FDTDGrid class describing a grid in a model.
-    """
-
-    if x == 0 or z == 0:
-        Iy = 0
-    else:
-        Iy = G.dx * (Hx[x, y, z] - Hx[x, y, z - 1]) + G.dz * (Hz[x - 1, y, z] - Hz[x, y, z])
-
-    return Iy
-
-
-def Iz(x, y, z, Hx, Hy, Hz, G):
-    """Calculates the z-component of current at a grid position.
-
-    Args:
-        x, y, z: floats for coordinates of position in grid.
-        Hx, Hy, Hz: numpy array of magnetic field values.
-        G: FDTDGrid class describing a grid in a model.
-    """
-
-    if x == 0 or y == 0:
-        Iz = 0
-    else:
-        Iz = G.dx * (Hx[x, y - 1, z] - Hx[x, y, z]) + G.dy * (Hy[x, y, z] - Hy[x - 1, y, z])
-
-    return Iz

gprMax/grid/fdtd_grid.py

@@ -555,6 +555,60 @@ class FDTDGrid:
         # binary representation of floating point number.
         self.dt = round_value(self.dt, decimalplaces=decimal.getcontext().prec - 1)
 
+    def calculate_Ix(self, x: int, y: int, z: int) -> float:
+        """Calculates the x-component of current at a grid position.
+
+        Args:
+            x: x coordinate of position in grid
+            y: y coordinate of position in grid
+            z: z coordinate of position in grid
+        """
+
+        if y == 0 or z == 0:
+            Ix = 0
+        else:
+            Ix = self.dy * (self.Hy[x, y, z - 1] - self.Hy[x, y, z]) + self.dz * (
+                self.Hz[x, y, z] - self.Hz[x, y - 1, z]
+            )
+
+        return Ix
+
+    def calculate_Iy(self, x: int, y: int, z: int) -> float:
+        """Calculates the y-component of current at a grid position.
+
+        Args:
+            x: x coordinate of position in grid
+            y: y coordinate of position in grid
+            z: z coordinate of position in grid
+        """
+
+        if x == 0 or z == 0:
+            Iy = 0
+        else:
+            Iy = self.dx * (self.Hx[x, y, z] - self.Hx[x, y, z - 1]) + self.dz * (
+                self.Hz[x - 1, y, z] - self.Hz[x, y, z]
+            )
+
+        return Iy
+
+    def calculate_Iz(self, x: int, y: int, z: int) -> float:
+        """Calculates the y-component of current at a grid position.
+
+        Args:
+            x: x coordinate of position in grid
+            y: y coordinate of position in grid
+            z: z coordinate of position in grid
+        """
+
+        if x == 0 or y == 0:
+            Iz = 0
+        else:
+            Iz = self.dx * (self.Hx[x, y - 1, z] - self.Hx[x, y, z]) + self.dy * (
+                self.Hy[x, y, z] - self.Hy[x - 1, y, z]
+            )
+
+        return Iz
+
 
 def dispersion_analysis(G):
     """Analysis of numerical dispersion (Taflove et al, 2005, p112) -

gprMax/sources.py

@@ -22,7 +22,6 @@ import numpy as np
 
 import gprMax.config as config
 
-from .fields_outputs import Ix, Iy, Iz
 from .utilities.utilities import round_value
 
 
@@ -32,14 +31,14 @@ class Source:
     def __init__(self):
         self.ID = None
         self.polarisation = None
-        self.xcoord = None
-        self.ycoord = None
-        self.zcoord = None
+        self.xcoord = 0
+        self.ycoord = 0
+        self.zcoord = 0
         self.xcoordorigin = None
         self.ycoordorigin = None
         self.zcoordorigin = None
-        self.start = None
-        self.stop = None
+        self.start = 0
+        self.stop = 0
         self.waveformID = None
 
     def calculate_waveform_values(self, G):
@@ -49,10 +48,14 @@ class Source:
             G: FDTDGrid class describing a grid in a model.
         """
         # Waveform values for sources that need to be calculated on whole timesteps
-        self.waveformvalues_wholedt = np.zeros((G.iterations), dtype=config.sim_config.dtypes["float_or_double"])
+        self.waveformvalues_wholedt = np.zeros(
+            (G.iterations), dtype=config.sim_config.dtypes["float_or_double"]
+        )
 
         # Waveform values for sources that need to be calculated on half timesteps
-        self.waveformvalues_halfdt = np.zeros((G.iterations), dtype=config.sim_config.dtypes["float_or_double"])
+        self.waveformvalues_halfdt = np.zeros(
+            (G.iterations), dtype=config.sim_config.dtypes["float_or_double"]
+        )
 
         waveform = next(x for x in G.waveforms if x.ID == self.waveformID)
 
@@ -63,7 +66,9 @@ class Source:
                 # delay in the start
                 time -= self.start
                 self.waveformvalues_wholedt[iteration] = waveform.calculate_value(time, G.dt)
-                self.waveformvalues_halfdt[iteration] = waveform.calculate_value(time + 0.5 * G.dt, G.dt)
+                self.waveformvalues_halfdt[iteration] = waveform.calculate_value(
+                    time + 0.5 * G.dt, G.dt
+                )
 
 
 class VoltageSource(Source):
@@ -274,7 +279,9 @@ def htod_src_arrays(sources, G, queue=None):
 
     srcinfo1 = np.zeros((len(sources), 4), dtype=np.int32)
     srcinfo2 = np.zeros((len(sources)), dtype=config.sim_config.dtypes["float_or_double"])
-    srcwaves = np.zeros((len(sources), G.iterations), dtype=config.sim_config.dtypes["float_or_double"])
+    srcwaves = np.zeros(
+        (len(sources), G.iterations), dtype=config.sim_config.dtypes["float_or_double"]
+    )
     for i, src in enumerate(sources):
         srcinfo1[i, 0] = src.xcoord
         srcinfo1[i, 1] = src.ycoord
@@ -399,11 +406,15 @@ class TransmissionLine(Source):
 
         # Update all the voltage values along the line
         self.voltage[1 : self.nl] -= (
-            self.resistance * (config.c * G.dt / self.dl) * (self.current[1 : self.nl] - self.current[0 : self.nl - 1])
+            self.resistance
+            * (config.c * G.dt / self.dl)
+            * (self.current[1 : self.nl] - self.current[0 : self.nl - 1])
         )
 
         # Update the voltage at the position of the one-way injector excitation
-        self.voltage[self.srcpos] += (config.c * G.dt / self.dl) * self.waveformvalues_halfdt[iteration]
+        self.voltage[self.srcpos] += (config.c * G.dt / self.dl) * self.waveformvalues_halfdt[
+            iteration
+        ]
 
         # Update ABC before updating current
         self.update_abc(G)
@@ -425,7 +436,9 @@ class TransmissionLine(Source):
 
         # Update the current one cell before the position of the one-way injector excitation
         self.current[self.srcpos - 1] += (
-            (1 / self.resistance) * (config.c * G.dt / self.dl) * self.waveformvalues_wholedt[iteration]
+            (1 / self.resistance)
+            * (config.c * G.dt / self.dl)
+            * self.waveformvalues_wholedt[iteration]
         )
 
     def update_electric(self, iteration, updatecoeffsE, ID, Ex, Ey, Ez, G):
@@ -458,6 +471,7 @@ class TransmissionLine(Source):
             elif self.polarisation == "z":
                 Ez[i, j, k] = -self.voltage[self.antpos] / G.dz
 
+    # TODO: Add type information (if can avoid circular dependency)
     def update_magnetic(self, iteration, updatecoeffsH, ID, Hx, Hy, Hz, G):
         """Updates current value in transmission line from magnetic field values
             in the main grid.
@@ -478,12 +492,12 @@ class TransmissionLine(Source):
             k = self.zcoord
 
             if self.polarisation == "x":
-                self.current[self.antpos] = Ix(i, j, k, G.Hx, G.Hy, G.Hz, G)
+                self.current[self.antpos] = G.calculate_Ix(i, j, k)
 
             elif self.polarisation == "y":
-                self.current[self.antpos] = Iy(i, j, k, G.Hx, G.Hy, G.Hz, G)
+                self.current[self.antpos] = G.calculate_Iy(i, j, k)
 
             elif self.polarisation == "z":
-                self.current[self.antpos] = Iz(i, j, k, G.Hx, G.Hy, G.Hz, G)
+                self.current[self.antpos] = G.calculate_Iz(i, j, k)
 
             self.update_current(iteration, G)

tests/grid/test_fdtd_grid.py

@@ -0,0 +1,121 @@
+from typing import Callable
+
+import numpy as np
+import pytest
+from numpy.testing import assert_allclose, assert_array_equal
+from pytest import param
+
+from gprMax.grid.fdtd_grid import FDTDGrid
+
+
+def get_current_in_3d_grid(
+    get_current_func: Callable[[int, int, int], float], shape: tuple[int, ...]
+) -> np.ndarray:
+    """Helper function to get current as a 3D grid
+
+    Args:
+        get_current_func: Function that returns the current value at a
+            point on a grid.
+        shape: shape of the grid
+
+    Returns:
+        result: 3D grid containing current values
+    """
+    result = np.empty(shape)
+    for i in range(shape[0]):
+        for j in range(shape[1]):
+            for k in range(shape[2]):
+                result[i, j, k] = get_current_func(i, j, k)
+    return result
+
+
+def tile(a: float, b: float, n: int = 1) -> np.ndarray:
+    """Helper function to tile two numbers
+
+    Args:
+        a: first value
+        b: second value
+        n: number of repititions
+
+    Returns:
+        c: tiled numpy array
+    """
+    return np.tile([[[a, b], [b, a]], [[b, a], [a, b]]], (n, n, n))
+
+
+@pytest.mark.parametrize(
+    "dy,dz,Hy,Hz,expected",
+    [
+        (0, 0, (1, 3), (0.5, 0.8), 0),
+        (0.1, 0.5, (0, 0), (0.5, 0.8), -0.15),
+        (0.1, 0.5, (1, 3), (0, 0), 0.2),
+        (0.1, 0, (1, 3), (0.5, 0.8), 0.2),
+        (0, 0.5, (1, 3), (0.5, 0.8), -0.15),
+        (0.1, 0.5, (1, 3), (0.5, 0.8), 0.05),
+    ],
+)
+def test_calculate_Ix(dy, dz, Hy, Hz, expected, size=2):
+    grid = FDTDGrid()
+    grid.dy = dy
+    grid.dz = dz
+    grid.Hy = tile(Hy[0], Hy[1], size)
+    grid.Hz = tile(Hz[0], Hz[1], size)
+
+    actual_current = get_current_in_3d_grid(grid.calculate_Ix, grid.Hy.shape)
+    expected_current = tile(expected, -expected, size)
+    expected_current[:, 0, :] = 0
+    expected_current[:, :, 0] = 0
+
+    assert_allclose(actual_current, expected_current)
+
+
+@pytest.mark.parametrize(
+    "dx,dz,Hx,Hz,expected",
+    [
+        (0, 0, (1, 3), (0.5, 0.8), 0),
+        (0.1, 0.5, (0, 0), (0.5, 0.8), -0.15),
+        (0.1, 0.5, (1, 3), (0, 0), 0.2),
+        (0.1, 0, (1, 3), (0.5, 0.8), 0.2),
+        (0, 0.5, (1, 3), (0.5, 0.8), -0.15),
+        (0.1, 0.5, (1, 3), (0.5, 0.8), 0.05),
+    ],
+)
+def test_calculate_Iy(dx, dz, Hx, Hz, expected, size=2):
+    grid = FDTDGrid()
+    grid.dx = dx
+    grid.dz = dz
+    grid.Hx = tile(Hx[0], Hx[1], size)
+    grid.Hz = tile(Hz[0], Hz[1], size)
+
+    actual_current = get_current_in_3d_grid(grid.calculate_Iy, grid.Hx.shape)
+    expected_current = tile(-expected, expected, size)
+    expected_current[0, :, :] = 0
+    expected_current[:, :, 0] = 0
+
+    assert_allclose(actual_current, expected_current)
+
+
+@pytest.mark.parametrize(
+    "dx,dy,Hx,Hy,expected",
+    [
+        (0, 0, (1, 3), (0.5, 0.8), 0),
+        (0.1, 0.5, (0, 0), (0.5, 0.8), -0.15),
+        (0.1, 0.5, (1, 3), (0, 0), 0.2),
+        (0.1, 0, (1, 3), (0.5, 0.8), 0.2),
+        (0, 0.5, (1, 3), (0.5, 0.8), -0.15),
+        (0.1, 0.5, (1, 3), (0.5, 0.8), 0.05),
+    ],
+)
+def test_calculate_Iz(dx, dy, Hx, Hy, expected, size=2):
+    grid = FDTDGrid()
+    grid.dx = dx
+    grid.dy = dy
+    grid.Hx = tile(Hx[0], Hx[1], size)
+    grid.Hy = tile(Hy[0], Hy[1], size)
+
+    actual_current = get_current_in_3d_grid(grid.calculate_Iz, grid.Hx.shape)
+    expected_current = tile(expected, -expected, size)
+    expected_current[0, :, :] = 0
+    expected_current[:, 0, :] = 0
+
+    assert_allclose(actual_current, expected_current)