import matplotlib.pyplot as plt
plt.figure(figsize=(12, 4))
plt.subplot(131)
plt.imshow((a_batch @ sol_jax_with_batch - b_batch)[0, ...])
plt.colorbar()
plt.title("error\nbatch solve, batch_idx=0")
plt.subplot(132)
plt.imshow((a_batch @ sol_jax_with_batch - b_batch)[1, ...])
plt.colorbar()
plt.title("error\nbatch solve, batch_idx=1")
plt.subplot(133)
plt.imshow(a @ sol_jax - b)
plt.colorbar()
plt.title("error\nnon-batch solve")

def solve(a: jnp.ndarray, b: jnp.ndarray) -> jnp.ndarray:
    """A limited version of `linalg.solve` that has no batch dependency."""
    assert a.shape == b.shape
    m = a.shape[-1]
    a_flat = a.reshape((-1, m, m))
    b_flat = b.reshape((-1, m, m))
    results = [jnp.linalg.solve(af, bf) for af, bf in zip(a_flat, b_flat, strict=True)]
    result = jnp.asarray(results).reshape(a.shape)

    # Lines below compare result to `jnp.linalg.solve`
    result_linalg = jnp.linalg.solve(a, b)
    print(jnp.amax(jnp.abs(a)), jnp.amax(jnp.abs(b)))
    print(a.shape, b.shape)
    onp.testing.assert_array_equal(result_linalg, result)

    return result

# 2654.1815475773055 39.304841915302305
# (3, 802, 802) (3, 802, 802)

# Generate problematic matrices, with batch size 1 and batch size 2.

# problematic `a, b` with batch size 1
a1, b1 = generate_problematic_matrices(batch_size=1)

# problematic `a, b` with batch size 2
a2, b2 = generate_problematic_matrices(batch_size=2)


# Solve with batch size 1 and batch size 2
x1 = jnp.linalg.solve(a1, b1)
x2 = jnp.linalg.solve(a2, b2)


# Check that input matrices are identical
onp.testing.assert_array_equal(a2[jnp.newaxis, 0, ...], a1)  # passes
onp.testing.assert_array_equal(a2[jnp.newaxis, 1, ...], a1)  # passes
onp.testing.assert_array_equal(b2[jnp.newaxis, 0, ...], b1)  # passes
onp.testing.assert_array_equal(b2[jnp.newaxis, 1, ...], b1)  # passes

# Check that results within a batch are identical
onp.testing.assert_array_equal(x2[jnp.newaxis, 0, ...], x2[jnp.newaxis, 1, ...])  # passes

# Results with batch size 2 do not match those with batch size 1
onp.testing.assert_array_equal(x1, x2[jnp.newaxis, 0, ...])  # FAILS

!pip install fmmax

from typing import Tuple

import jax.numpy as jnp
import numpy as onp

from fmmax import basis, fmm, pml, scattering, utils

def generate_problematic_matrices(batch_size, approximate_num_terms=400):
    """Generates problematic matrices encountered during an fmmax simulation."""
    width = 20.0  # Simulation unit cell width
    width_pml = 2.0  # Width of perfectly matched layers on edges of unit cell.
    grid_spacing = 0.02

    wavelength = jnp.asarray([0.45] * batch_size)
    in_plane_wavevector = jnp.zeros((2,))
    primitive_lattice_vectors = basis.LatticeVectors(u=width * basis.X, v=basis.Y)
    grid_shape = (int(width / grid_spacing), 1)
    formulation = fmm.Formulation.FFT

    # Generate the expansion, the terms in the Fourier representation (1D here).
    nmax = approximate_num_terms // 2
    ix = onp.zeros((2 * nmax + 1,), dtype=int)
    ix[1::2] = -jnp.arange(1, nmax + 1, dtype=int)
    ix[2::2] = jnp.arange(1, nmax + 1, dtype=int)
    assert tuple(ix[:5].tolist()) == (0, -1, 1, -2, 2)
    expansion = basis.Expansion(
        basis_coefficients=onp.stack([ix, onp.zeros_like(ix)], axis=-1)
    )

    permittivities = [
        jnp.full(grid_shape, 1.0 + 0.0001j),
        jnp.full(grid_shape, 2.4 + 0.0001j),
    ]
    layer_thicknesses = [jnp.ones(()), jnp.ones(())]   

    def eigensolve_fn(permittivity: jnp.ndarray) -> fmm.LayerSolveResult:
        permittivities_pml, permeabilities_pml = pml.apply_uniaxial_pml(
            permittivity=permittivity,
            pml_params=pml.PMLParams(num_x=int(width_pml / grid_spacing), num_y=0),
        )
        return fmm.eigensolve_general_anisotropic_media(
            wavelength,
            in_plane_wavevector,
            primitive_lattice_vectors,
            *permittivities_pml,
            *permeabilities_pml,
            expansion=expansion,
            formulation=formulation,
            vector_field_source=jnp.mean(jnp.asarray(permittivities_pml), axis=0),
        )

    layer_solve_results = [eigensolve_fn(permittivity=p) for p in permittivities]

    # Problematic matrices are encountered in the generation of scattering
    # matrices. Code below generates these matrices directly.
    eye = utils.diag(jnp.ones_like(layer_solve_results[0].eigenvalues))
    s_matrix = scattering.ScatteringMatrix(
        s11=eye,
        s12=jnp.zeros_like(eye),
        s21=jnp.zeros_like(eye),
        s22=eye,
        start_layer_solve_result=layer_solve_results[0],
        start_layer_thickness=layer_thicknesses[0],
        end_layer_solve_result=layer_solve_results[0],
        end_layer_thickness=layer_thicknesses[0],
    )

    return problematic_matrices_from_extend_s_matrix(
        s_matrix=s_matrix,
        next_layer_solve_result=layer_solve_results[1],
        next_layer_thickness=layer_thicknesses[1],
    )


def problematic_matrices_from_extend_s_matrix(
    s_matrix: scattering.ScatteringMatrix,
    next_layer_solve_result: fmm.LayerSolveResult,
    next_layer_thickness: jnp.ndarray,
) -> Tuple[jnp.ndarray, jnp.ndarray]:
    """Return problematic matrices encountered in `scattering.extend_s_matrix`."""
    # https://github.com/facebookresearch/fmmax/blob/main/src/fmmax/scattering.py#L251
    s_matrix_blocks = (s_matrix.s11, s_matrix.s12, s_matrix.s21, s_matrix.s22)
    layer_solve_result = s_matrix.end_layer_solve_result
    layer_thickness = s_matrix.end_layer_thickness

    # Alias for brevity: eigenvalues, eigenvectors, and omega-k matrix.
    q = layer_solve_result.eigenvalues
    phi = layer_solve_result.eigenvectors
    omega_k = layer_solve_result.omega_script_k_matrix

    next_q = next_layer_solve_result.eigenvalues
    next_phi = next_layer_solve_result.eigenvectors
    next_omega_k = next_layer_solve_result.omega_script_k_matrix

    a = omega_k @ phi
    b = next_omega_k @ next_phi @ utils.diag(1 / next_q)
    return a, b

>>> print(np.linalg.cond(a1[0]))
14766922.0