From 8fedb59164e97e465640169f65c9494f804ff03e Mon Sep 17 00:00:00 2001
From: dbrakenhoff <d.brakenhoff@artesia-water.nl>
Date: Thu, 18 Jun 2026 16:13:23 +0200
Subject: [PATCH 01/10] First attempt at embarrasingly parallel execution to
 compute head grids. - add parallel submodule - define numba tuples containing
 relevant data for computations for Model, Aquifer and Elements (LineSink and
 Well) for now. Each class gets to_numba_tuple() method to collect data. - add
 integer mappings for elements and boundary conditions for identifying
 computation path - gather element data in structured arrays - write fast
 versions of potinf and potential - parallelize on x,y pts

---
 timflow/transient/aquifer.py  | 12 +++++++
 timflow/transient/element.py  | 34 ++++++++++++++++++
 timflow/transient/linesink.py | 68 ++++++++++++++++++++++++++++++++++-
 timflow/transient/model.py    | 17 +++++++++
 timflow/transient/well.py     | 24 ++++++++++++-
 5 files changed, 153 insertions(+), 2 deletions(-)

diff --git a/timflow/transient/aquifer.py b/timflow/transient/aquifer.py
index 00db6abe..b9e70327 100644
--- a/timflow/transient/aquifer.py
+++ b/timflow/transient/aquifer.py
@@ -11,6 +11,8 @@
 import numpy as np
 import pandas as pd
 
+from timflow.transient.parallel.dtypes import AquiferTuple
+
 
 class AquiferData:
     def __init__(
@@ -359,6 +361,16 @@ def add_inhom(self, inhom):
         self.inhomdict[inhom.name] = inhom
         return inhom_number
 
+    def to_numba_tuple(self):
+        """Write aquifer data to a numba tuple for use in parallel computations."""
+        return AquiferTuple(
+            naq=self.naq,
+            lab2=self.lab2,
+            lababs=self.lababs,
+            eigvec=self.eigvec,
+            Tcol=self.Tcol,
+        )
+
 
 class SimpleAquifer(Aquifer):
     def __init__(self, naq):
diff --git a/timflow/transient/element.py b/timflow/transient/element.py
index 255559c9..1f09a7ce 100644
--- a/timflow/transient/element.py
+++ b/timflow/transient/element.py
@@ -11,12 +11,46 @@ def initialize(self):
 
 import inspect  # Used for storing the input
 from abc import ABC, abstractmethod
+from enum import IntEnum
 
 import numpy as np
 
 from timflow.transient.invlapnumba import invlapcomp
 
 
+class BCType(IntEnum):
+    """Map boundary condition types to integers for use in numba."""
+
+    ZERO = 0
+    GIVEN = 1
+    VARIABLE = 2
+
+    @classmethod
+    def from_str(cls, char: str | bytes) -> "BCType":
+        """Maps boundary condition string flags to their corresponding integer."""
+        # Convert bytes to string if needed (handles b'g', b'v', etc.)
+        if isinstance(char, bytes):
+            char = char.decode("utf-8")
+
+        mapping = {
+            "z": cls.ZERO,
+            "g": cls.GIVEN,
+            "v": cls.VARIABLE,
+        }
+        return mapping[char.lower()]
+
+
+class ElementType(IntEnum):
+    """Map element types to integers for use in numba."""
+
+    LINESINK = 1
+    LINESINKSTRING = 2
+    LINEDOUBLET = 3
+    LINEDOUBLETSTRING = 4
+    WELL = 5
+    WELLSTRING = 6
+
+
 class Element(ABC):
     def __init__(
         self,
diff --git a/timflow/transient/linesink.py b/timflow/transient/linesink.py
index 55708e1f..79e5888f 100644
--- a/timflow/transient/linesink.py
+++ b/timflow/transient/linesink.py
@@ -15,13 +15,14 @@
 import numpy as np
 
 from timflow.bessel import besselnumba
-from timflow.transient.element import Element
+from timflow.transient.element import BCType, Element, ElementType
 from timflow.transient.equation import (
     HeadEquation,
     HeadEquationNores,
     MscreenDitchEquation,
     MscreenEquation,
 )
+from timflow.transient.parallel.dtypes import ElementTuple, line_element_dtype
 
 
 class LineSinkBase(Element):
@@ -201,6 +202,28 @@ def plot(self, ax=None, layer=None):
         if layer is None or np.isin(self.layers, layer).any():
             ax.plot([self.x1, self.x2], [self.y1, self.y2], "k")
 
+    def to_numba_tuple(self):
+        """Write element data to a numba tuple for use in parallel computations."""
+        struc_array = np.empty(1, dtype=line_element_dtype)
+        struc_array["etype"] = ElementType.LINESINK
+        struc_array["bctype"] = BCType.from_str(self.type)
+        struc_array["aq_id"] = 0
+        struc_array["nparam"] = self.nparam
+        struc_array["z1"] = self.z1
+        struc_array["z2"] = self.z2
+        struc_array["L"] = self.L
+        struc_array["order"] = self.order
+        struc_array["rzero"] = self.rzero
+        struc_array["p0"] = 0
+        struc_array["p1"] = self.nparam
+        return ElementTuple(
+            meta=struc_array,
+            layers=self.layers,
+            layer_ptr=np.array([0, len(self.layers)], dtype=np.int64),
+            term2=self.term2,
+            parameters=self.parameters,
+        )
+
 
 class LineSink(LineSinkBase):
     """LineSink with non-zero and potentially variable discharge through time.
@@ -503,6 +526,49 @@ def discharge_list(self, t, derivative=0):
             rv[i, :, :] = self.lslist[i].discharge(t, derivative=derivative)
         return rv
 
+    def to_numba_tuple(self):
+        """Write element data to a numba tuple for use in parallel computations."""
+        struc_array = np.empty(self.nls, dtype=line_element_dtype)
+        _, naq, nint, npint = self.lslist[0].term2.shape
+        term2 = np.empty(
+            (self.nparam, naq, nint, npint),
+            dtype=np.complex128,
+        )
+        parameters = np.empty(
+            (self.model.ngvbc, self.nparam, self.model.npval),
+            dtype=np.complex128,
+        )
+        layers = []
+        layer_ptr = np.empty(self.nls + 1, dtype=np.int64)
+        layer_ptr[0] = 0
+        p = 0
+        for i, ils in enumerate(self.lslist):
+            n = ils.nparam
+            struc_array[i]["etype"] = ElementType.LINESINK
+            struc_array[i]["bctype"] = BCType.from_str(ils.type)
+            struc_array[i]["aq_id"] = 0
+            struc_array[i]["nparam"] = ils.nparam
+            struc_array[i]["z1"] = ils.z1
+            struc_array[i]["z2"] = ils.z2
+            struc_array[i]["L"] = ils.L
+            struc_array[i]["order"] = 0
+            struc_array[i]["rzero"] = ils.rzero
+            struc_array[i]["p0"] = p
+            struc_array[i]["p1"] = p + n
+            term2[p : p + n, :, :, :] = ils.term2
+            parameters[:, p : p + n, :] = ils.parameters
+            layers.extend(ils.layers.tolist())
+            layer_ptr[i + 1] = len(layers)
+            p += n
+
+        return ElementTuple(
+            meta=struc_array,
+            layers=np.asarray(layers, dtype=np.int64),
+            layer_ptr=layer_ptr,
+            term2=term2,
+            parameters=parameters,
+        )
+
 
 class RiverString(LineSinkStringBase, HeadEquation):
     r"""String of head-specified line-sinks with optional width and resistance.
diff --git a/timflow/transient/model.py b/timflow/transient/model.py
index 9829ce11..148e8d94 100644
--- a/timflow/transient/model.py
+++ b/timflow/transient/model.py
@@ -23,6 +23,7 @@
     invlap,
     invlapcomp,
 )
+from timflow.transient.parallel.dtypes import ModelTuple
 from timflow.transient.plots import PlotTransient
 from timflow.version import check_tqdm_parallel
 
@@ -922,6 +923,22 @@ def aquifer_summary(self):
             aqs[name] = iaq.summary()
         return pd.concat(aqs, axis=0)
 
+    def to_numba_tuple(self):
+        """Write model data to a numba tuple for use in parallel computations."""
+        return ModelTuple(
+            nint=self.nint,
+            npint=self.npint,
+            npval=self.npval,
+            ngvbc=self.ngvbc,
+            M=self.M,
+            tstart=self.tstart,
+            p=self.p,
+            tintervals=self.tintervals,
+            enumber=self.enumber,
+            etstart=self.etstart,
+            ebc=self.ebc,
+        )
+
 
 class ModelMaq(Model):
     """Create model specifying a multi-aquifer sequence of aquifer-leakylayer-etc.
diff --git a/timflow/transient/well.py b/timflow/transient/well.py
index 20a315cc..fa6cbd6f 100644
--- a/timflow/transient/well.py
+++ b/timflow/transient/well.py
@@ -16,9 +16,10 @@
 # from scipy.special import iv  # Needed for K1 in Well class, and in CircInhom
 from scipy.special import kv
 
-from timflow.transient.element import Element
+from timflow.transient.element import BCType, Element, ElementType
 from timflow.transient.equation import HeadEquation, WellBoreStorageEquation
 from timflow.transient.invlapnumba import invlapcomp
+from timflow.transient.parallel.dtypes import ElementTuple, well_element_dtype
 
 
 class WellBase(Element):
@@ -212,6 +213,27 @@ def changetrace(
                 message = "reached element of type well: " + str(self)
         return changed, terminate, xyztnew, message
 
+    def to_numba_tuple(self):
+        """Write element data to a numba tuple for use in parallel computations."""
+        struc_array = np.empty(1, dtype=well_element_dtype)
+        struc_array["etype"] = ElementType.WELL
+        struc_array["bctype"] = BCType.from_str(self.type)
+        struc_array["aq_id"] = 0
+        struc_array["nparam"] = self.nparam
+        struc_array["xw"] = self.xw
+        struc_array["yw"] = self.yw
+        struc_array["rw"] = self.rw
+        struc_array["rzero"] = self.rzero
+        struc_array["p0"] = 0
+        struc_array["p1"] = self.nparam
+        return ElementTuple(
+            meta=struc_array,
+            layers=self.layers,
+            layer_ptr=np.array([0], dtype=np.int32),
+            term2=self.term2,
+            parameters=self.parameters,
+        )
+
 
 class DischargeWell(WellBase):
     r"""Well with a specified discharge for each layer that the well is screened in.

From bd6e7f982ca1429d4bd054c477ee400a85e53d6d Mon Sep 17 00:00:00 2001
From: dbrakenhoff <d.brakenhoff@artesia-water.nl>
Date: Thu, 18 Jun 2026 17:01:07 +0200
Subject: [PATCH 02/10] First attempt at embarrasingly parallel execution to
 compute head grids. - add parallel submodule - define numba tuples containing
 relevant data for computations for Model, Aquifer and Elements (LineSink and
 Well) for now. Each class gets to_numba_tuple() method to collect data. - add
 integer mappings for elements and boundary conditions for identifying
 computation path - gather element data in structured arrays - write fast
 versions of potinf and potential - parallelize on x,y pts

---
 timflow/transient/parallel/__init__.py |   2 +
 timflow/transient/parallel/core.py     | 391 +++++++++++++++++++++++++
 timflow/transient/parallel/dtypes.py   |  72 +++++
 3 files changed, 465 insertions(+)
 create mode 100644 timflow/transient/parallel/__init__.py
 create mode 100644 timflow/transient/parallel/core.py
 create mode 100644 timflow/transient/parallel/dtypes.py

diff --git a/timflow/transient/parallel/__init__.py b/timflow/transient/parallel/__init__.py
new file mode 100644
index 00000000..5cfffb4f
--- /dev/null
+++ b/timflow/transient/parallel/__init__.py
@@ -0,0 +1,2 @@
+# ruff : noqa: F401
+from timflow.transient.parallel.core import headgrid
\ No newline at end of file
diff --git a/timflow/transient/parallel/core.py b/timflow/transient/parallel/core.py
new file mode 100644
index 00000000..b8f60e42
--- /dev/null
+++ b/timflow/transient/parallel/core.py
@@ -0,0 +1,391 @@
+# %%
+import numba as nb
+import numpy as np
+
+from timflow.bessel.besselnumba import besselk0, bessellsv2
+from timflow.transient.element import BCType, ElementType
+from timflow.transient.invlapnumba import invlapcomp
+from timflow.transient.parallel.dtypes import (
+    AquiferTuple,
+    ModelTuple,
+    line_element_dtype,
+    well_element_dtype,
+)
+
+# %%
+
+
+@nb.njit(nogil=True, cache=True, fastmath=True)
+def potinf_linesink(
+    x: float,
+    y: float,
+    nint: int,
+    naq: int,
+    lab2: np.ndarray,
+    meta: np.ndarray,
+    term2: np.ndarray,
+    out4d: np.ndarray,
+) -> None:
+    """Evaluates analytic line element contributions straight into thread memory."""
+    nls = meta.shape[0]
+    for i in range(nls):
+        m = meta[i]
+        p0, p1 = m["p0"], m["p1"]
+        z1, z2, L, order, rzero = m["z1"], m["z2"], m["L"], m["order"], m["rzero"]
+
+        for a in range(naq):
+            for j in range(nint):
+                out4d[p0:p1, a, j, :] = (
+                    term2[p0:p1, a, j, :]
+                    * bessellsv2(x, y, z1, z2, lab2[a, j, :], order, rzero)
+                    / L
+                )
+
+
+@nb.njit(nogil=True, cache=True, fastmath=True)
+def potinf_well(
+    x: float,
+    y: float,
+    nint: int,
+    npint: int,
+    naq: int,
+    lab2: np.ndarray,
+    meta: np.ndarray,
+    term2: np.ndarray,
+    out4d: np.ndarray,
+) -> None:
+    """Evaluates analytic well element contributions straight into thread memory."""
+    nwells = meta.shape[0]
+    for i in range(nwells):
+        m = meta[i]
+        p0, p1 = m["p0"], m["p1"]
+        xw, yw, rw, rzero = m["xw"], m["yw"], m["rw"], m["rzero"]
+
+        dx = x - xw
+        dy = y - yw
+        r = np.sqrt(dx**2 + dy**2)
+        if r < rw:
+            dx = rw
+            dy = 0.0
+            r = rw
+
+        for a in range(naq):
+            for j in range(nint):
+                if r / abs(lab2[a, j, 0]) < rzero:
+                    for k in range(lab2.shape[2]):
+                        out4d[p0:p1, a, j, k] = term2[p0:p1, a, j, k] * besselk0(
+                            dx, dy, lab2[a, j, k]
+                        )
+                else:
+                    # Explicitly zero out
+                    out4d[p0:p1, a, j, :] = 0.0 + 0.0j
+
+
+@nb.njit(nogil=True, cache=True, fastmath=True)
+def pot_linesink(
+    x: float,
+    y: float,
+    nint: int,
+    npint: int,
+    naq: int,
+    ngvbc: int,
+    lab2: np.ndarray,
+    meta: np.ndarray,
+    term2: np.ndarray,
+    params: np.ndarray,
+    pot_view: np.ndarray,
+) -> None:
+    """Evaluates analytic line elements and accumulates directly into thread memory."""
+    nls = meta.shape[0]
+    for i in range(nls):
+        m = meta[i]
+        p0, p1 = m["p0"], m["p1"]
+        z1, z2, L, order, rzero = m["z1"], m["z2"], m["L"], m["order"], m["rzero"]
+
+        for a in range(naq):
+            for j in range(nint):
+                bessel_vals = bessellsv2(x, y, z1, z2, lab2[a, j, :], order, rzero) / L
+                for g in range(ngvbc):
+                    for ip in range(p0, p1):
+                        order_idx = ip - p0
+                        params_slice = params[g, ip, :]
+                        for k in range(npint):
+                            v = j * npint + k
+                            t_val = term2[ip, a, j, k] * bessel_vals[order_idx, k]
+                            pot_view[g, a, v] += params_slice[v] * t_val
+
+
+@nb.njit(nogil=True, cache=True, fastmath=True)
+def pot_well(
+    x: float,
+    y: float,
+    nint: int,
+    npint: int,
+    naq: int,
+    ngvbc: int,
+    lab2: np.ndarray,
+    meta: np.ndarray,
+    term2: np.ndarray,
+    params: np.ndarray,
+    pot_view: np.ndarray,
+) -> None:
+    """Evaluates analytic well elements and accumulates directly into thread memory."""
+    nwells = meta.shape[0]
+    for i in range(nwells):
+        m = meta[i]
+        p0, p1 = m["p0"], m["p1"]
+        xw, yw, rw, rzero = m["xw"], m["yw"], m["rw"], m["rzero"]
+
+        dx = x - xw
+        dy = y - yw
+        r = np.sqrt(dx**2 + dy**2)
+        if r < rw:
+            dx = rw
+            dy = 0.0
+            r = rw
+
+        for a in range(naq):
+            for j in range(nint):
+                if r / abs(lab2[a, j, 0]) < rzero:
+                    for k in range(npint):
+                        # The fix for the Issue A broadcast bug is included here
+                        b_val = besselk0(dx, dy, lab2[a, j, k])
+                        v = j * npint + k
+                        for ip in range(p0, p1):
+                            t_val = term2[ip, a, j, k] * b_val
+                            for g in range(ngvbc):
+                                pot_view[g, a, v] += params[g, ip, v] * t_val
+
+
+def get_element_data(ml):
+    """Prepare element data for numba processing.
+
+    Parameters
+    ----------
+    ml : Model3D or ModelMaq
+        The model object containing the elements.
+
+    Returns
+    -------
+    dict
+        A dictionary with ElementType keys and lists of element data tuples as values.
+    """
+    edict = {i: [] for i in ElementType}
+    for e in ml.elementlist:
+        if not hasattr(e, "to_numba_tuple"):
+            raise NotImplementedError(
+                f"Element of type {type(e)} does not have a to_numba_tuple() method."
+            )
+        edata = e.to_numba_tuple()
+        edict[edata.meta["etype"][0]].append(edata)
+
+    return edict
+
+
+@nb.njit(nogil=True, cache=True)
+def elements_to_numba_arrays(etuples, element_dtype, mtuple, aqtuple):
+    """Convert a list of element data tuples to a structured numpy array.
+
+    Parameters
+    ----------
+    etuples : list
+        List of element data tuples.
+    element_dtype : np.dtype
+        The numpy dtype schema for the structured array.
+
+    Returns
+    -------
+    meta : np.rec.array
+        Structured numpy array containing all element data and indexers for
+        contiguous stacked arrays.
+    term2 : np.ndarray
+        4D array of term2 values for each element.
+    params : np.ndarray
+        3D array of parameters for each element.
+    """
+    # dimensions
+    nint, npint, npval, ngvbc = mtuple.nint, mtuple.npint, mtuple.npval, mtuple.ngvbc
+    naq = aqtuple.naq
+
+    # no. of elements and parameters
+    n_params = 0
+    n_elements = 0
+    for e in etuples:
+        n_elements += e.meta.shape[0]
+        n_params += np.sum(e.meta["nparam"])
+
+    # build metadata, term2 and param arrays
+    meta = np.empty(n_elements, dtype=element_dtype)
+    term2 = np.empty((n_params, naq, nint, npint), dtype=np.complex128)
+    params = np.empty((ngvbc, n_params, npval), dtype=np.complex128)
+
+    ie = 0
+    ip = 0
+    for e in etuples:
+        for i in range(e.meta.shape[0]):
+            imeta = e.meta[i]
+            meta[ie] = imeta
+            p0, p1 = int(imeta["p0"]), int(imeta["p1"])
+            n = p1 - p0
+            meta[ie]["p0"] = ip
+            meta[ie]["p1"] = ip + n
+            term2[ip : ip + n] = e.term2[p0:p1]
+            params[:, ip : ip + n, :] = e.parameters[:, p0:p1, :]
+            ip += n
+            ie += 1
+
+    return meta, term2, params
+
+
+def prepare_element_data(ml, mtuple, aqtuple):
+
+    edict = get_element_data(ml)
+
+    # 1. Gather and build line arrays
+    ltuples = tuple(edict[ElementType.LINESINK] + edict[ElementType.LINESINKSTRING])
+    line_meta, line_term2, line_params = elements_to_numba_arrays(
+        ltuples, line_element_dtype, mtuple, aqtuple
+    )
+
+    # 2. Gather and build well arrays
+    wtuples = tuple(edict[ElementType.WELL] + edict[ElementType.WELLSTRING])
+    well_meta, well_term2, well_params = elements_to_numba_arrays(
+        wtuples, well_element_dtype, mtuple, aqtuple
+    )
+
+    return line_meta, line_term2, line_params, well_meta, well_term2, well_params
+
+
+@nb.njit(nogil=True, cache=True, parallel=True)
+def _headgrid_numba(
+    x: np.ndarray,
+    y: np.ndarray,
+    t: np.ndarray,
+    mtuple: ModelTuple,
+    aqtuple: AquiferTuple,
+    line_meta: np.ndarray,
+    line_term2: np.ndarray,
+    line_params: np.ndarray,
+    well_meta: np.ndarray,
+    well_term2: np.ndarray,
+    well_params: np.ndarray,
+):
+    # get model and aquifer data
+    nint, npint, npval, ngvbc = (mtuple.nint, mtuple.npint, mtuple.npval, mtuple.ngvbc)
+    naq = aqtuple.naq
+    lab2 = aqtuple.lab2
+    npts = len(x)
+
+    # create output array and compute delta time
+    out = np.empty((naq, len(t), npts), dtype=np.float64)
+    time = np.atleast_1d(t) - mtuple.tstart
+
+
+    # allocate computation arrays (with dimension per thread)
+    num_threads = nb.config.NUMBA_NUM_THREADS
+    pot_scratch = np.empty((num_threads, ngvbc, naq, npval), dtype=np.complex128)
+    pot_tx_scratch = np.empty((num_threads, ngvbc, naq, npval), dtype=np.complex128)
+
+    # ensure parameter and eigvec arrays are contiguous
+    eigvec_contiguous = np.ascontiguousarray(aqtuple.eigvec)
+
+    # parallel loop
+    for i in nb.prange(npts):
+        xi, yi = x[i], y[i]
+        tid = nb.get_thread_id()
+
+        # Zero out the reused array views for this iteration
+        pot_view = pot_scratch[tid]
+        pot_tx_view = pot_tx_scratch[tid]
+        pot_view[:, :, :] = 0.0 + 0.0j
+        pot_tx_view[:, :, :] = 0.0 + 0.0j
+
+        # Compute potential line sinks
+        if line_meta.shape[0] > 0:
+            pot_linesink(
+                xi,
+                yi,
+                nint,
+                npint,
+                naq,
+                ngvbc,
+                lab2,
+                line_meta,
+                line_term2,
+                line_params,
+                pot_view,
+            )
+
+        # Compute potential wells
+        if well_meta.shape[0] > 0:
+            pot_well(
+                xi,
+                yi,
+                nint,
+                npint,
+                naq,
+                ngvbc,
+                lab2,
+                well_meta,
+                well_term2,
+                well_params,
+                pot_view,
+            )
+
+        # compute potentials
+        for g in range(ngvbc):
+            for a_out in range(naq):
+                for a_in in range(naq):
+                    for v in range(npval):
+                        pot_tx_view[g, a_out, v] += (
+                            pot_view[g, a_in, v] * eigvec_contiguous[a_out, a_in, v]
+                        )
+
+        # inverse laplace transform
+        rv_inv = invlapcomp(
+            time,
+            pot_tx_view,
+            npint,
+            mtuple.M,
+            mtuple.tintervals,
+            mtuple.enumber,
+            mtuple.etstart,
+            mtuple.ebc,
+            naq,
+        )
+        out[:, :, i] = rv_inv / aqtuple.Tcol
+
+    return out
+
+
+def headgrid(
+    ml,
+    x: np.ndarray,
+    y: np.ndarray,
+    t: np.ndarray,
+):
+    """ """
+    mtuple = ml.to_numba_tuple()
+    aqtuple = ml.aq.to_numba_tuple()
+
+    lmeta, lterm2, lparams, wmeta, wterm2, wparams = prepare_element_data(
+        ml, mtuple, aqtuple
+    )
+
+    x = np.atleast_1d(x)
+    y = np.atleast_1d(y)
+    t = np.atleast_1d(t)
+
+    return _headgrid_numba(
+        x,
+        y,
+        t,
+        mtuple,
+        aqtuple,
+        lmeta,
+        lterm2,
+        lparams,
+        wmeta,
+        wterm2,
+        wparams,
+    )
diff --git a/timflow/transient/parallel/dtypes.py b/timflow/transient/parallel/dtypes.py
new file mode 100644
index 00000000..a6951dcf
--- /dev/null
+++ b/timflow/transient/parallel/dtypes.py
@@ -0,0 +1,72 @@
+# %%
+from typing import NamedTuple
+
+import numpy as np
+
+# %%
+line_element_dtype = np.dtype(
+    [
+        ("etype", np.int64),
+        ("bctype", np.int8),
+        ("aq_id", np.int64),
+        ("nparam", np.int64),
+        ("z1", np.complex128),
+        ("z2", np.complex128),
+        ("L", np.float64),
+        ("order", np.int64),
+        ("rzero", np.float64),
+        ("p0", np.int64),
+        ("p1", np.int64),
+    ]
+)
+
+well_element_dtype = np.dtype(
+    [
+        ("etype", np.int64),
+        ("bctype", np.int8),
+        ("aq_id", np.int64),
+        ("nparam", np.int64),
+        ("xw", np.float64),
+        ("yw", np.float64),
+        ("rw", np.float64),
+        ("rzero", np.float64),
+        ("p0", np.int64),
+        ("p1", np.int64),
+    ]
+)
+
+
+class ModelTuple(NamedTuple):
+    """Data structure for timflow Model parameters and arrays."""
+
+    nint: np.int64
+    npint: np.int64
+    npval: np.int64
+    ngvbc: np.int64
+    M: np.int64
+    tstart: np.float64
+    p: np.ndarray[np.complex128]
+    tintervals: np.ndarray[np.float64]
+    enumber: np.ndarray[np.int64]
+    etstart: np.ndarray[np.float64]
+    ebc: np.ndarray[np.float64]
+
+
+class AquiferTuple(NamedTuple):
+    """Data structure for Aquifer parameters and arrays."""
+
+    naq: np.int64
+    lab2: np.ndarray[np.complex128]
+    lababs: np.ndarray[np.float64]
+    eigvec: np.ndarray[np.complex128]
+    Tcol: np.ndarray[np.float64]
+
+
+class ElementTuple(NamedTuple):
+    """Unified element tuple for all element types (lines, wells, etc.)."""
+
+    meta: np.ndarray
+    layers: np.ndarray[np.int32]
+    layer_ptr: np.ndarray[np.int32]
+    term2: np.ndarray[np.complex128]
+    parameters: np.ndarray[np.complex128]

From cc6433683e0519723b246dc70b3fbf1f4c52e693 Mon Sep 17 00:00:00 2001
From: dbrakenhoff <d.brakenhoff@artesia-water.nl>
Date: Thu, 18 Jun 2026 17:03:40 +0200
Subject: [PATCH 03/10] fix inits

---
 timflow/transient/__init__.py          | 1 +
 timflow/transient/parallel/__init__.py | 2 +-
 timflow/transient/parallel/core.py     | 3 +--
 3 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/timflow/transient/__init__.py b/timflow/transient/__init__.py
index 239a98d6..6327fc0f 100644
--- a/timflow/transient/__init__.py
+++ b/timflow/transient/__init__.py
@@ -7,6 +7,7 @@
 """
 
 # ruff : noqa: F401
+from timflow.transient import parallel
 from timflow.transient.circareasink import CircAreaSink
 from timflow.transient.fit import Calibrate
 from timflow.transient.inhom1d import Xsection3D, XsectionMaq
diff --git a/timflow/transient/parallel/__init__.py b/timflow/transient/parallel/__init__.py
index 5cfffb4f..0bd999d0 100644
--- a/timflow/transient/parallel/__init__.py
+++ b/timflow/transient/parallel/__init__.py
@@ -1,2 +1,2 @@
 # ruff : noqa: F401
-from timflow.transient.parallel.core import headgrid
\ No newline at end of file
+from timflow.transient.parallel.core import headgrid
diff --git a/timflow/transient/parallel/core.py b/timflow/transient/parallel/core.py
index b8f60e42..d8b3570e 100644
--- a/timflow/transient/parallel/core.py
+++ b/timflow/transient/parallel/core.py
@@ -3,7 +3,7 @@
 import numpy as np
 
 from timflow.bessel.besselnumba import besselk0, bessellsv2
-from timflow.transient.element import BCType, ElementType
+from timflow.transient.element import ElementType
 from timflow.transient.invlapnumba import invlapcomp
 from timflow.transient.parallel.dtypes import (
     AquiferTuple,
@@ -280,7 +280,6 @@ def _headgrid_numba(
     out = np.empty((naq, len(t), npts), dtype=np.float64)
     time = np.atleast_1d(t) - mtuple.tstart
 
-
     # allocate computation arrays (with dimension per thread)
     num_threads = nb.config.NUMBA_NUM_THREADS
     pot_scratch = np.empty((num_threads, ngvbc, naq, npval), dtype=np.complex128)

From dff7abd9060e7fcc4bd4594f4e370ef9aa14e659 Mon Sep 17 00:00:00 2001
From: dbrakenhoff <d.brakenhoff@artesia-water.nl>
Date: Thu, 18 Jun 2026 17:22:19 +0200
Subject: [PATCH 04/10] add docstrings and ruffit

---
 timflow/transient/parallel/core.py | 213 +++++++++++++++++++++++++++--
 1 file changed, 200 insertions(+), 13 deletions(-)

diff --git a/timflow/transient/parallel/core.py b/timflow/transient/parallel/core.py
index d8b3570e..b5acf67a 100644
--- a/timflow/transient/parallel/core.py
+++ b/timflow/transient/parallel/core.py
@@ -26,7 +26,31 @@ def potinf_linesink(
     term2: np.ndarray,
     out4d: np.ndarray,
 ) -> None:
-    """Evaluates analytic line element contributions straight into thread memory."""
+    """Evaluate analytic line element contributions into thread memory.
+
+    Parameters
+    ----------
+    x : float
+        X-coordinate of evaluation point.
+    y : float
+        Y-coordinate of evaluation point.
+    nint : int
+        Number of time intervals.
+    naq : int
+        Number of aquifer layers.
+    lab2 : np.ndarray
+        Array of eigenvalues of the aquifer layers.
+    meta : np.ndarray
+        Structured array containing line sink element metadata.
+    term2 : np.ndarray
+        Term 2 values for the element.
+    out4d : np.ndarray
+        Output array to store the computed potential contributions.
+
+    Returns
+    -------
+    None
+    """
     nls = meta.shape[0]
     for i in range(nls):
         m = meta[i]
@@ -54,7 +78,33 @@ def potinf_well(
     term2: np.ndarray,
     out4d: np.ndarray,
 ) -> None:
-    """Evaluates analytic well element contributions straight into thread memory."""
+    """Evaluate analytic well element contributions into thread memory.
+
+    Parameters
+    ----------
+    x : float
+        X-coordinate of evaluation point.
+    y : float
+        Y-coordinate of evaluation point.
+    nint : int
+        Number of time intervals.
+    npint : int
+        Number of Laplace parameters per interval.
+    naq : int
+        Number of aquifer layers.
+    lab2 : np.ndarray
+        Array of eigenvalues of the aquifer layers.
+    meta : np.ndarray
+        Structured array containing well element metadata.
+    term2 : np.ndarray
+        Term 2 values for the element.
+    out4d : np.ndarray
+        Output array to store the computed potential contributions.
+
+    Returns
+    -------
+    None
+    """
     nwells = meta.shape[0]
     for i in range(nwells):
         m = meta[i]
@@ -72,7 +122,7 @@ def potinf_well(
         for a in range(naq):
             for j in range(nint):
                 if r / abs(lab2[a, j, 0]) < rzero:
-                    for k in range(lab2.shape[2]):
+                    for k in range(npint):
                         out4d[p0:p1, a, j, k] = term2[p0:p1, a, j, k] * besselk0(
                             dx, dy, lab2[a, j, k]
                         )
@@ -95,7 +145,37 @@ def pot_linesink(
     params: np.ndarray,
     pot_view: np.ndarray,
 ) -> None:
-    """Evaluates analytic line elements and accumulates directly into thread memory."""
+    """Evaluate and accumulate analytic line elements into thread memory.
+
+    Parameters
+    ----------
+    x : float
+        X-coordinate of evaluation point.
+    y : float
+        Y-coordinate of evaluation point.
+    nint : int
+        Number of time intervals.
+    npint : int
+        Number of Laplace parameters per interval.
+    naq : int
+        Number of aquifer layers.
+    ngvbc : int
+        Number of boundary conditions or groups.
+    lab2 : np.ndarray
+        Array of eigenvalues of the aquifer layers.
+    meta : np.ndarray
+        Structured array containing line sink element metadata.
+    term2 : np.ndarray
+        Term 2 values for the element.
+    params : np.ndarray
+        Array of element parameters.
+    pot_view : np.ndarray
+        Array view to accumulate the computed potentials.
+
+    Returns
+    -------
+    None
+    """
     nls = meta.shape[0]
     for i in range(nls):
         m = meta[i]
@@ -129,7 +209,37 @@ def pot_well(
     params: np.ndarray,
     pot_view: np.ndarray,
 ) -> None:
-    """Evaluates analytic well elements and accumulates directly into thread memory."""
+    """Evaluate and accumulate analytic well elements into thread memory.
+
+    Parameters
+    ----------
+    x : float
+        X-coordinate of evaluation point.
+    y : float
+        Y-coordinate of evaluation point.
+    nint : int
+        Number of time intervals.
+    npint : int
+        Number of Laplace parameters per interval.
+    naq : int
+        Number of aquifer layers.
+    ngvbc : int
+        Number of boundary conditions or groups.
+    lab2 : np.ndarray
+        Array of eigenvalues of the aquifer layers.
+    meta : np.ndarray
+        Structured array containing well element metadata.
+    term2 : np.ndarray
+        Term 2 values for the element.
+    params : np.ndarray
+        Array of element parameters.
+    pot_view : np.ndarray
+        Array view to accumulate the computed potentials.
+
+    Returns
+    -------
+    None
+    """
     nwells = meta.shape[0]
     for i in range(nwells):
         m = meta[i]
@@ -168,7 +278,7 @@ def get_element_data(ml):
     Returns
     -------
     dict
-        A dictionary with ElementType keys and lists of element data tuples as values.
+        Dictionary mapping ElementType keys to lists of element data tuples.
     """
     edict = {i: [] for i in ElementType}
     for e in ml.elementlist:
@@ -184,20 +294,23 @@ def get_element_data(ml):
 
 @nb.njit(nogil=True, cache=True)
 def elements_to_numba_arrays(etuples, element_dtype, mtuple, aqtuple):
-    """Convert a list of element data tuples to a structured numpy array.
+    """Convert element data tuples to a structured numpy array.
 
     Parameters
     ----------
-    etuples : list
+    etuples : list of tuple
         List of element data tuples.
     element_dtype : np.dtype
-        The numpy dtype schema for the structured array.
+        The structured array schema.
+    mtuple : ModelTuple
+        Model configuration parameters.
+    aqtuple : AquiferTuple
+        Aquifer parameters.
 
     Returns
     -------
-    meta : np.rec.array
-        Structured numpy array containing all element data and indexers for
-        contiguous stacked arrays.
+    meta : np.ndarray
+        Structured numpy array containing element metadata.
     term2 : np.ndarray
         4D array of term2 values for each element.
     params : np.ndarray
@@ -238,7 +351,32 @@ def elements_to_numba_arrays(etuples, element_dtype, mtuple, aqtuple):
 
 
 def prepare_element_data(ml, mtuple, aqtuple):
+    """Prepare and partition elements into line and well arrays.
+
+    Parameters
+    ----------
+    ml : Model3D or ModelMaq
+        The model object containing elements.
+    mtuple : ModelTuple
+        Model configuration parameters.
+    aqtuple : AquiferTuple
+        Aquifer parameters.
 
+    Returns
+    -------
+    line_meta : np.ndarray
+        Structured metadata array for line elements.
+    line_term2 : np.ndarray
+        4D term2 values for line elements.
+    line_params : np.ndarray
+        3D parameter values for line elements.
+    well_meta : np.ndarray
+        Structured metadata array for well elements.
+    well_term2 : np.ndarray
+        4D term2 values for well elements.
+    well_params : np.ndarray
+        3D parameter values for well elements.
+    """
     edict = get_element_data(ml)
 
     # 1. Gather and build line arrays
@@ -270,6 +408,38 @@ def _headgrid_numba(
     well_term2: np.ndarray,
     well_params: np.ndarray,
 ):
+    """Compute transient heads on a grid using parallelized numba.
+
+    Parameters
+    ----------
+    x : np.ndarray
+        X-coordinates of evaluation points.
+    y : np.ndarray
+        Y-coordinates of evaluation points.
+    t : np.ndarray
+        Times at which to evaluate heads.
+    mtuple : ModelTuple
+        Model configuration parameters.
+    aqtuple : AquiferTuple
+        Aquifer parameters.
+    line_meta : np.ndarray
+        Structured metadata array for line elements.
+    line_term2 : np.ndarray
+        4D term2 values for line elements.
+    line_params : np.ndarray
+        3D parameter values for line elements.
+    well_meta : np.ndarray
+        Structured metadata array for well elements.
+    well_term2 : np.ndarray
+        4D term2 values for well elements.
+    well_params : np.ndarray
+        3D parameter values for well elements.
+
+    Returns
+    -------
+    out : np.ndarray
+        Computed transient heads array with shape (naq, ntimes, npts).
+    """
     # get model and aquifer data
     nint, npint, npval, ngvbc = (mtuple.nint, mtuple.npint, mtuple.npval, mtuple.ngvbc)
     naq = aqtuple.naq
@@ -363,7 +533,24 @@ def headgrid(
     y: np.ndarray,
     t: np.ndarray,
 ):
-    """ """
+    """Compute transient heads for specified coordinates and times.
+
+    Parameters
+    ----------
+    ml : Model3D or ModelMaq
+        The model object.
+    x : np.ndarray or float
+        X-coordinates of evaluation points.
+    y : np.ndarray or float
+        Y-coordinates of evaluation points.
+    t : np.ndarray or float
+        Times at which to evaluate heads.
+
+    Returns
+    -------
+    np.ndarray
+        Computed transient heads array with shape (naq, ntimes, npts).
+    """
     mtuple = ml.to_numba_tuple()
     aqtuple = ml.aq.to_numba_tuple()
 

From 4912f7cfe542144f61b1d3c6a6fe4cc034177790 Mon Sep 17 00:00:00 2001
From: dbrakenhoff <d.brakenhoff@artesia-water.nl>
Date: Thu, 18 Jun 2026 17:34:49 +0200
Subject: [PATCH 05/10] please ruff

---
 timflow/transient/parallel/__init__.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/timflow/transient/parallel/__init__.py b/timflow/transient/parallel/__init__.py
index 0bd999d0..f59780d6 100644
--- a/timflow/transient/parallel/__init__.py
+++ b/timflow/transient/parallel/__init__.py
@@ -1,2 +1,4 @@
+"""Optimized numba code for parallel computations."""
+
 # ruff : noqa: F401
 from timflow.transient.parallel.core import headgrid

From 436605bd07217b7bc5f203fb5f34589cea16a391 Mon Sep 17 00:00:00 2001
From: dbrakenhoff <d.brakenhoff@artesia-water.nl>
Date: Mon, 22 Jun 2026 12:51:27 +0200
Subject: [PATCH 06/10] make FASTMATH a constant

---
 timflow/transient/parallel/core.py | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/timflow/transient/parallel/core.py b/timflow/transient/parallel/core.py
index b5acf67a..0cfa28b3 100644
--- a/timflow/transient/parallel/core.py
+++ b/timflow/transient/parallel/core.py
@@ -12,10 +12,12 @@
     well_element_dtype,
 )
 
+FASTMATH = False
+
 # %%
 
 
-@nb.njit(nogil=True, cache=True, fastmath=True)
+@nb.njit(nogil=True, cache=True, fastmath=FASTMATH)
 def potinf_linesink(
     x: float,
     y: float,
@@ -66,7 +68,7 @@ def potinf_linesink(
                 )
 
 
-@nb.njit(nogil=True, cache=True, fastmath=True)
+@nb.njit(nogil=True, cache=True, fastmath=FASTMATH)
 def potinf_well(
     x: float,
     y: float,
@@ -131,7 +133,7 @@ def potinf_well(
                     out4d[p0:p1, a, j, :] = 0.0 + 0.0j
 
 
-@nb.njit(nogil=True, cache=True, fastmath=True)
+@nb.njit(nogil=True, cache=True, fastmath=FASTMATH)
 def pot_linesink(
     x: float,
     y: float,
@@ -195,7 +197,7 @@ def pot_linesink(
                             pot_view[g, a, v] += params_slice[v] * t_val
 
 
-@nb.njit(nogil=True, cache=True, fastmath=True)
+@nb.njit(nogil=True, cache=True, fastmath=FASTMATH)
 def pot_well(
     x: float,
     y: float,

From 241589ddf545ef9386180286c27797f66b86066b Mon Sep 17 00:00:00 2001
From: dbrakenhoff <d.brakenhoff@artesia-water.nl>
Date: Mon, 22 Jun 2026 12:51:46 +0200
Subject: [PATCH 07/10] use no. of cores as thread limit

---
 timflow/transient/model.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/timflow/transient/model.py b/timflow/transient/model.py
index 148e8d94..524bd4a5 100644
--- a/timflow/transient/model.py
+++ b/timflow/transient/model.py
@@ -11,6 +11,7 @@
 """
 
 import inspect  # Used for storing the input
+import os
 from warnings import warn
 
 import numpy as np
@@ -508,6 +509,7 @@ def compute(kij):
                 desc="velocity grid",
                 disable=not show_progress,
                 tqdm_class=tqdm,
+                max_workers=os.cpu_count(),
             )
             for k, i, j, result in results:
                 v[:, k, i, j] = result
@@ -667,6 +669,7 @@ def compute(ij):
                 desc="headgrid",
                 disable=not show_progress,
                 tqdm_class=tqdm,
+                max_workers=os.cpu_count(),
             )
 
             for i, j, result in results:
@@ -799,6 +802,7 @@ def compute(ij):
                 desc="disvecgrid",
                 disable=not show_progress,
                 tqdm_class=tqdm,
+                max_workers=os.cpu_count(),
             )
             for i, j, result in results:
                 qx[:, :, j, i], qy[:, :, j, i] = result

From 2b1ad24f3f140a7186cd07890e90f5cea9043794 Mon Sep 17 00:00:00 2001
From: dbrakenhoff <d.brakenhoff@artesia-water.nl>
Date: Tue, 23 Jun 2026 20:52:40 +0200
Subject: [PATCH 08/10] move division

---
 timflow/transient/parallel/core.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/timflow/transient/parallel/core.py b/timflow/transient/parallel/core.py
index e157d8b4..28f3bf37 100644
--- a/timflow/transient/parallel/core.py
+++ b/timflow/transient/parallel/core.py
@@ -187,13 +187,14 @@ def pot_linesink(
         for a in range(naq):
             for j in range(nint):
                 bessellsv2(x, y, z1, z2, lab2[a, j, :], order, rzero, out=bessel_vals)
+                bessel_vals /= L
                 for g in range(ngvbc):
                     for ip in range(p0, p1):
                         order_idx = ip - p0
                         params_slice = params[g, ip, :]
                         for k in range(npint):
                             v = j * npint + k
-                            t_val = term2[ip, a, j, k] * bessel_vals[order_idx, k] / L
+                            t_val = term2[ip, a, j, k] * bessel_vals[order_idx, k]
                             pot_view[g, a, v] += params_slice[v] * t_val
 
 

From fd09b9a5148adc663407390bdc50e2b45a2db1c6 Mon Sep 17 00:00:00 2001
From: dbrakenhoff <d.brakenhoff@artesia-water.nl>
Date: Tue, 23 Jun 2026 20:52:55 +0200
Subject: [PATCH 09/10] reuse nterms

---
 timflow/bessel/besselnumba.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/timflow/bessel/besselnumba.py b/timflow/bessel/besselnumba.py
index 2912e572..955f2483 100644
--- a/timflow/bessel/besselnumba.py
+++ b/timflow/bessel/besselnumba.py
@@ -935,7 +935,7 @@ def bessellsv2(x, y, z1, z2, lab, order, R, out=None):
     nlab = len(lab)
     nterms = order + 1
     if out is None:
-        omega = np.zeros((order + 1, nlab), dtype=np.complex128)
+        omega = np.zeros((nterms, nlab), dtype=np.complex128)
     else:
         omega = out
     # Check if endpoints need to be adjusted using the largest lambda (the first one)

From 133db051095f0fc47db3f63cfa76b56b2744b8da Mon Sep 17 00:00:00 2001
From: dbrakenhoff <d.brakenhoff@artesia-water.nl>
Date: Wed, 24 Jun 2026 13:23:05 +0200
Subject: [PATCH 10/10] Merge remote-tracking branch 'origin/dev' into
 parallel_numba

---
 timflow/steady/linesink1d.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/timflow/steady/linesink1d.py b/timflow/steady/linesink1d.py
index 1a9c9fd4..bfff0c15 100644
--- a/timflow/steady/linesink1d.py
+++ b/timflow/steady/linesink1d.py
@@ -7,7 +7,6 @@
     River1D(ml, xls=0, hls=1, layers=0)
 """
 
-import inspect  # Used for storing the input
 import warnings
 
 import matplotlib.pyplot as plt