Models

class BGKSim(LBMBase):
    """
    BGK simulation class.

    This class implements the Bhatnagar-Gross-Krook (BGK) approximation for the collision step in the Lattice Boltzmann Method.
    """

    def __init__(self, **kwargs):
        super().__init__(**kwargs)

    @partial(jit, static_argnums=(0,))
    def collision(self, f):
        """
        BGK collision step for lattice.

        The collision step is where the main physics of the LBM is applied. In the BGK approximation,
        the distribution function is relaxed towards the equilibrium distribution function.
        """
        f = self.precision_policy.cast_to_compute(f)
        rho, u = self.update_macroscopic(f)
        feq = self.equilibrium(rho, u, cast_output=False)
        fneq = f - feq
        fout = f - self.omega * fneq
        if self.force is not None:
            fout = self.apply_force(fout, feq, rho, u)
        return self.precision_policy.cast_to_output(fout)

class KBCSim(LBMBase):
    """
    KBC simulation class.

    This class implements the Karlin-Bösch-Chikatamarla (KBC) model for the collision step in the Lattice Boltzmann Method.
    """

    def __init__(self, **kwargs):
        if kwargs.get("lattice").name != "D3Q27" and kwargs.get("nz") > 0:
            raise ValueError("KBC collision operator in 3D must only be used with D3Q27 lattice.")
        super().__init__(**kwargs)

    @partial(jit, static_argnums=(0,))
    def collision(self, f):
        """
        KBC collision step for lattice.
        """
        f = self.precision_policy.cast_to_compute(f)
        tiny = 1e-32
        beta = self.omega * 0.5
        rho, u = self.update_macroscopic(f)
        feq = self.equilibrium(rho, u, cast_output=False)
        fneq = f - feq
        if self.dim == 2:
            deltaS = self.fdecompose_shear_d2q9(fneq) * rho / 4.0
        else:
            deltaS = self.fdecompose_shear_d3q27(fneq) * rho
        deltaH = fneq - deltaS
        invBeta = 1.0 / beta
        gamma = invBeta - (2.0 - invBeta) * self.entropic_scalar_product(deltaS, deltaH, feq) / (
            tiny + self.entropic_scalar_product(deltaH, deltaH, feq)
        )

        fout = f - beta * (2.0 * deltaS + gamma[..., None] * deltaH)

        # add external force
        if self.force is not None:
            fout = self.apply_force(fout, feq, rho, u)
        return self.precision_policy.cast_to_output(fout)

    @partial(jit, static_argnums=(0,))
    def collision_modified(self, f):
        """
        Alternative KBC collision step for lattice.
        Note:
        At low Reynolds number the orignal KBC collision above produces inaccurate results because
        it does not check for the entropy increase/decrease. The KBC stabalizations should only be
        applied in principle to cells whose entropy decrease after a regular BGK collision. This is
        the case in most cells at higher Reynolds numbers and hence a check may not be needed.
        Overall the following alternative collision is more reliable and may replace the original
        implementation. The issue at the moment is that it is about 60-80% slower than the above method.
        """
        f = self.precision_policy.cast_to_compute(f)
        tiny = 1e-32
        beta = self.omega * 0.5
        rho, u = self.update_macroscopic(f)
        feq = self.equilibrium(rho, u, castOutput=False)

        # Alternative KBC: only stabalizes for voxels whose entropy decreases after BGK collision.
        f_bgk = f - self.omega * (f - feq)
        H_fin = jnp.sum(f * jnp.log(f / self.w), axis=-1, keepdims=True)
        H_fout = jnp.sum(f_bgk * jnp.log(f_bgk / self.w), axis=-1, keepdims=True)

        # the rest is identical to collision_deprecated
        fneq = f - feq
        if self.dim == 2:
            deltaS = self.fdecompose_shear_d2q9(fneq) * rho / 4.0
        else:
            deltaS = self.fdecompose_shear_d3q27(fneq) * rho
        deltaH = fneq - deltaS
        invBeta = 1.0 / beta
        gamma = invBeta - (2.0 - invBeta) * self.entropic_scalar_product(deltaS, deltaH, feq) / (
            tiny + self.entropic_scalar_product(deltaH, deltaH, feq)
        )

        f_kbc = f - beta * (2.0 * deltaS + gamma[..., None] * deltaH)
        fout = jnp.where(H_fout > H_fin, f_kbc, f_bgk)

        # add external force
        if self.force is not None:
            fout = self.apply_force(fout, feq, rho, u)
        return self.precision_policy.cast_to_output(fout)

    @partial(jit, static_argnums=(0,), inline=True)
    def entropic_scalar_product(self, x, y, feq):
        """
        Compute the entropic scalar product of x and y to approximate gamma in KBC.

        Returns
        -------
        jax.numpy.array
            Entropic scalar product of x, y, and feq.
        """
        return jnp.sum(x * y / feq, axis=-1)

    @partial(jit, static_argnums=(0,), inline=True)
    def fdecompose_shear_d2q9(self, fneq):
        """
        Decompose fneq into shear components for D2Q9 lattice.

        Parameters
        ----------
        fneq (jax.numpy.array): Non-equilibrium distribution function.

        Returns
        -------
        jax.numpy.array
            Shear components of fneq.
        """
        Pi = self.momentum_flux(fneq)
        N = Pi[..., 0] - Pi[..., 2]
        s = jnp.zeros_like(fneq)
        s = s.at[..., 6].set(N)
        s = s.at[..., 3].set(N)
        s = s.at[..., 2].set(-N)
        s = s.at[..., 1].set(-N)
        s = s.at[..., 8].set(Pi[..., 1])
        s = s.at[..., 4].set(-Pi[..., 1])
        s = s.at[..., 5].set(-Pi[..., 1])
        s = s.at[..., 7].set(Pi[..., 1])

        return s

    @partial(jit, static_argnums=(0,), inline=True)
    def fdecompose_shear_d3q27(self, fneq):
        """
        Decompose fneq into shear components for D3Q27 lattice.

        Parameters
        ----------
        fneq (jax.numpy.ndarray): Non-equilibrium distribution function.

        Returns
        -------
        jax.numpy.ndarray
            Shear components of fneq.
        """
        # if self.grid.dim == 3:
        #     diagonal    = (0, 3, 5)
        #     offdiagonal = (1, 2, 4)
        # elif self.grid.dim == 2:
        #     diagonal    = (0, 2)
        #     offdiagonal = (1,)

        # c=
        # array([[0, 0, 0],-----0
        #        [0, 0, -1],----1
        #        [0, 0, 1],-----2
        #        [0, -1, 0],----3
        #        [0, -1, -1],---4
        #        [0, -1, 1],----5
        #        [0, 1, 0],-----6
        #        [0, 1, -1],----7
        #        [0, 1, 1],-----8
        #        [-1, 0, 0],----9
        #        [-1, 0, -1],--10
        #        [-1, 0, 1],---11
        #        [-1, -1, 0],--12
        #        [-1, -1, -1],-13
        #        [-1, -1, 1],--14
        #        [-1, 1, 0],---15
        #        [-1, 1, -1],--16
        #        [-1, 1, 1],---17
        #        [1, 0, 0],----18
        #        [1, 0, -1],---19
        #        [1, 0, 1],----20
        #        [1, -1, 0],---21
        #        [1, -1, -1],--22
        #        [1, -1, 1],---23
        #        [1, 1, 0],----24
        #        [1, 1, -1],---25
        #        [1, 1, 1]])---26
        Pi = self.momentum_flux(fneq)
        Nxz = Pi[..., 0] - Pi[..., 5]
        Nyz = Pi[..., 3] - Pi[..., 5]

        # For c = (i, 0, 0), c = (0, j, 0) and c = (0, 0, k)
        s = jnp.zeros_like(fneq)
        s = s.at[..., 9].set((2.0 * Nxz - Nyz) / 6.0)
        s = s.at[..., 18].set((2.0 * Nxz - Nyz) / 6.0)
        s = s.at[..., 3].set((-Nxz + 2.0 * Nyz) / 6.0)
        s = s.at[..., 6].set((-Nxz + 2.0 * Nyz) / 6.0)
        s = s.at[..., 1].set((-Nxz - Nyz) / 6.0)
        s = s.at[..., 2].set((-Nxz - Nyz) / 6.0)

        # For c = (i, j, 0)
        s = s.at[..., 12].set(Pi[..., 1] / 4.0)
        s = s.at[..., 24].set(Pi[..., 1] / 4.0)
        s = s.at[..., 21].set(-Pi[..., 1] / 4.0)
        s = s.at[..., 15].set(-Pi[..., 1] / 4.0)

        # For c = (i, 0, k)
        s = s.at[..., 10].set(Pi[..., 2] / 4.0)
        s = s.at[..., 20].set(Pi[..., 2] / 4.0)
        s = s.at[..., 19].set(-Pi[..., 2] / 4.0)
        s = s.at[..., 11].set(-Pi[..., 2] / 4.0)

        # For c = (0, j, k)
        s = s.at[..., 8].set(Pi[..., 4] / 4.0)
        s = s.at[..., 4].set(Pi[..., 4] / 4.0)
        s = s.at[..., 7].set(-Pi[..., 4] / 4.0)
        s = s.at[..., 5].set(-Pi[..., 4] / 4.0)

        return s

class AdvectionDiffusionBGK(LBMBase):
    """
    Advection Diffusion Model based on the BGK model.
    """

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.vel = kwargs.get("vel", None)
        if self.vel is None:
            raise ValueError("Velocity must be specified for AdvectionDiffusionBGK.")

    @partial(jit, static_argnums=(0,))
    def collision(self, f):
        """
        BGK collision step for lattice.
        """
        f = self.precision_policy.cast_to_compute(f)
        rho = jnp.sum(f, axis=-1, keepdims=True)
        feq = self.equilibrium(rho, self.vel, cast_output=False)
        fneq = f - feq
        fout = f - self.omega * fneq
        return self.precision_policy.cast_to_output(fout)

class MRTSim(LBMBase):
    """
    Multi-relaxation time model.
    """

    def __init__(self, **kwargs):
        kwargs.update({"omega": 1.0})
        super().__init__(**kwargs)
        self.s_rho = kwargs.get("s_rho")
        self.s_e = kwargs.get("s_e")
        self.s_eta = kwargs.get("s_eta")
        self.s_j = kwargs.get("s_j")
        self.s_q = kwargs.get("s_q")
        self.s_v = self.omega
        self.M_inv = jnp.array(
            np.transpose(np.linalg.inv(kwargs.get("M"))),
            dtype=self.precision_policy.compute_dtype,
        )
        self.M = jnp.array(np.transpose(kwargs.get("M")), dtype=self.precision_policy.compute_dtype)
        if isinstance(self.lattice, LatticeD2Q9):
            self.S = jnp.array(
                np.diag([self.s_rho, self.s_e, self.s_eta, self.s_j, self.s_q, self.s_j, self.s_q, self.s_v, self.s_v]),
                dtype=self.precision_policy.compute_dtype,
            )
        elif isinstance(self.lattice, LatticeD3Q19):
            self.s_pi = kwargs.get("s_pi")
            self.s_m = kwargs.get("s_m")
            self.S = jnp.array(
                np.diag([
                    self.s_rho,
                    self.s_e,
                    self.s_eta,
                    self.s_j,
                    self.s_q,
                    self.s_j,
                    self.s_q,
                    self.s_j,
                    self.s_q,
                    self.s_v,
                    self.s_pi,
                    self.s_v,
                    self.s_pi,
                    self.s_v,
                    self.s_v,
                    self.s_v,
                    self.s_m,
                    self.s_m,
                    self.s_m,
                ]),
                dtype=self.precision_policy.compute_dtype,
            )
        else:
            NotImplementedError(f"Lattice type {self.lattice.name} has not been implemented")

        # Fused collision matrix K = M @ S @ M_inv, replacing the three separate moment-space matrix multiplies
        # (f @ M, relax, @ M_inv) in collision() with symbolic, sparse-coefficient columns compiled from K's
        # nonzero entries - see collision() and MultiphaseMRT.collision (same fusion identity, without a
        # surface-tension term since this is the single-phase model).
        collision_matrix = np.asarray(jnp.dot(jnp.dot(self.M, self.S), self.M_inv))
        columns = []
        for output_direction in range(self.lattice.q):
            columns.append(
                tuple(
                    (input_direction, np.float32(collision_matrix[input_direction, output_direction]))
                    for input_direction in range(self.lattice.q)
                    if not np.isclose(collision_matrix[input_direction, output_direction], 0.0, atol=1e-7)
                )
            )
        self.collision_matrix = jnp.array(collision_matrix, dtype=self.precision_policy.compute_dtype)
        self.collision_terms = tuple(columns)

    @partial(jit, static_argnums=(0,))
    def collision(self, f):
        """
        MRT collision step for lattice, using a symbolic (sparse-coefficient) fused collision matrix instead of
        three separate moment-space matrix multiplies. See MultiphaseMRT.collision for the fusion identity.
        """
        f = self.precision_policy.cast_to_compute(f)
        rho, u = self.update_macroscopic(f)
        feq = self.equilibrium(rho, u)
        difference = f - feq

        outputs = []
        for output_direction, terms in enumerate(self.collision_terms):
            relaxed = sum(difference[..., input_direction] * coefficient for input_direction, coefficient in terms)
            outputs.append(f[..., output_direction] - relaxed)
        fout = jnp.stack(outputs, axis=-1)

        if self.force is not None:
            fout = fout + self._compute_force_delta_feq(rho, u)
        return self.precision_policy.cast_to_output(fout)

    @partial(jit, static_argnums=(0,), inline=True)
    def _compute_force_delta_feq(self, rho, u):
        """
        Real-space compact EDM difference delta_feq = feq(rho, u + du) - feq(rho, u), computed directly from
        cu, dcu and delta_usqr instead of building a full feq_force array via equilibrium(). Shared by
        collision (added directly, real space) and apply_force (transformed into moment space with M, for
        callers using the unfused, moment-space collision path).

        Parameters
        ----------
        rho (jax.numpy.ndarray): Density field.

        u (jax.numpy.ndarray): Velocity field.

        Returns
        -------
        (jax.numpy.ndarray): Real-space delta_feq.
        """
        du = self.get_force()
        c = jnp.array(self.c, dtype=self.precision_policy.compute_dtype)
        cu = 3.0 * jnp.dot(u, c)
        dcu = 3.0 * jnp.dot(du, c)
        delta_usqr = 1.5 * (2.0 * jnp.sum(u * du, axis=-1, keepdims=True) + jnp.sum(jnp.square(du), axis=-1, keepdims=True))
        return rho * self.w * (dcu * (1.0 + cu + 0.5 * dcu) - delta_usqr)

    @partial(jit, static_argnums=(0,), inline=True)
    def apply_force(self, m, meq, rho, u):
        """
        add force based on exact-difference method due to Kupershtokh, in moment space.

        LBMBase.apply_force subtracts feq in population space, which does not apply here since m/meq are
        moment-space quantities. Instead, the real-space delta_feq (see _compute_force_delta_feq) is
        transformed into moment space with M: since M is linear, dot(delta_feq, M) ==
        dot(feq(rho, u + du), M) - dot(feq(rho, u), M), an exact substitute for subtracting a moment-space
        equilibrium from a population-space one. Provided for callers using the unfused, moment-space collision
        path (e.g. Thermal.apply_force_thermal, when the wrapped fluid_solver is an MRTSim); collision() itself
        adds delta_feq directly in real space instead.

        Parameters
        ----------
        m (jax.numpy.ndarray): Post-collision moments.

        meq (jax.numpy.ndarray): Equilibrium moments. Unused - kept for interface compatibility, since the
            compact difference formula only needs rho, u and the force.

        rho (jax.numpy.ndarray): Density field.

        u (jax.numpy.ndarray): Velocity field.

        Returns
        -------
        (jax.numpy.ndarray): Post-collision moments with the force applied.

        References
        ----------
        1. Kupershtokh, A. (2004). New method of incorporating a body force term into the lattice Boltzmann
        equation. In Proceedings of the 5th International EHD Workshop (pp. 241-246). University of Poitiers.
        """
        delta_feq = self._compute_force_delta_feq(rho, u)
        return m + jnp.dot(delta_feq, self.M)