bf.geometry¶

def signed_distance(
    points: Array,
    a: Array,
    b: Array,
    poly_ids: np.ndarray,
    boundary: np.ndarray,
    chunk: int = 64,
    sign_tol: float = 1e-3,
) -> Array:
    """Signed distance (negative inside) from points to the union of polygons.

    Pure JAX, safe under ``jit``. Each polygon is filled by the nonzero winding
    rule (orientation does not matter) and the union distance is the minimum of
    the per-polygon signed distances, so overlapping polygons do not leave seams.

    The edges are streamed in chunks and every polygon is folded into the union
    as soon as its last edge has been seen, so memory is ~ (chunk + 6) * P no
    matter how many polygons there are; compute is ~ E * P.

    Where the nearest boundary point lies inside an edge (or the point lies on a
    vertex), the value is the signed distance to that edge's line (oriented with
    the polygon; averaged over tied edges), so the derivative stays correct on
    the boundary itself; near vertices it is sign * |p - v|.

    Args:
        points: (P, 2) query points.
        a: (E, 2) edge start points.
        b: (E, 2) edge end points.
        poly_ids: (E,) static polygon index of each edge (non-decreasing).
        boundary: (E,) static mask of the edges that count for the distance.
        chunk: edges processed per step.
        sign_tol: distance (points' units) from an edge's line within which the
            line's side is trusted over the winding number (round-off on the
            boundary); beyond it the winding number signs the distance.
    """
    if np.any(np.diff(poly_ids) < 0):
        raise ValueError("poly_ids must be sorted (edges of a polygon contiguous)")
    poly_ids = np.asarray(poly_ids)
    lay = _chunk_layout(poly_ids, np.asarray(boundary), chunk)
    # polygon orientation (+1 counter-clockwise), to sign the distance to edge lines
    num_polygons = int(poly_ids.max()) + 1 if poly_ids.size else 0
    cross = jax.lax.stop_gradient(a[:, 0] * b[:, 1] - b[:, 0] * a[:, 1])
    orient = jnp.sign(jax.ops.segment_sum(cross, jnp.asarray(poly_ids), num_segments=num_polygons))
    orient = orient[jnp.asarray(poly_ids)][lay["edge_index"]]  # (steps, C)
    a = a[lay["edge_index"]]  # (steps, C, 2)
    b = b[lay["edge_index"]]
    segments = chunk + 1  # local ids 0..chunk-1, plus one discarded segment
    k = jnp.arange(segments)[:, None]
    px = points[None, :, 0]
    py = points[None, :, 1]
    inf = jnp.asarray(jnp.inf, dtype=points.dtype)

    def seg_sum(x: Array, ids: Array) -> Array:
        return jax.ops.segment_sum(x, ids, num_segments=segments)

    @jax.checkpoint
    def body(carry: tuple[Array, ...], inp: tuple[Array, ...]) -> tuple[Any, None]:
        sd, open_m, open_lin, open_ni, open_n, open_w = carry
        ac, bc, oc, wid, did, n_local, cont_in, cont_out = inp
        ab = bc - ac
        ap = points[None, :, :] - ac[:, None, :]  # (C, P, 2)
        denom = jnp.sum(ab * ab, axis=-1)[:, None]
        t = jnp.sum(ap * ab[:, None, :], axis=-1) / jnp.where(denom > 0, denom, 1.0)
        d = ap - jnp.clip(t, 0.0, 1.0)[..., None] * ab[:, None, :]
        d2e = jnp.sum(d * d, axis=-1)  # (C, P)
        # signed distance to the edge line (negative inside): smooth across the
        # edge, unlike sqrt(d2), whose derivative vanishes on the boundary
        length = jnp.sqrt(jnp.where(denom > 0, denom, 1.0))
        cr = ab[:, None, 0] * ap[..., 1] - ab[:, None, 1] * ap[..., 0]
        lin = -oc[:, None] * cr / length
        # a point on a vertex is on both adjacent edges: their lines give the
        # derivative there (|p - v| has none at p = v)
        interior = ((t > 0) & (t < 1)) | (jax.lax.stop_gradient(d2e) <= 1e-12)

        # nearest edge per polygon: min over this chunk, merged with the open polygon.
        # The minimum only selects edges (no gradient); the differentiable value is
        # rebuilt from the selected edges with sums. segment_min's own derivative
        # rules are 0/0 for empty segments, which poisons derivatives of
        # derivatives (e.g. reverse mode through the normals' forward mode).
        d2_val = jax.lax.stop_gradient(d2e)
        open_val = jax.lax.stop_gradient(open_m)
        m_val = jax.ops.segment_min(d2_val, did, num_segments=segments)
        m_val = m_val.at[0].set(jnp.where(cont_in, jnp.minimum(m_val[0], open_val), m_val[0]))
        sel = (did < chunk)[:, None] & (d2_val <= m_val[did])
        sel_i = sel & interior
        lin_sum = seg_sum(jnp.where(sel_i, lin, 0.0), did)
        d2_sum = seg_sum(jnp.where(sel, d2e, 0.0), did)
        n_int = seg_sum(sel_i.astype(jnp.int32), did)
        n_sel = seg_sum(sel.astype(jnp.int32), did)
        take_open = cont_in & (open_val <= m_val[0])
        lin_sum = lin_sum.at[0].add(jnp.where(take_open, open_lin, 0.0))
        d2_sum = d2_sum.at[0].add(jnp.where(take_open & (open_n > 0), open_m, 0.0) * open_n)
        n_int = n_int.at[0].add(jnp.where(take_open, open_ni, 0))
        n_sel = n_sel.at[0].add(jnp.where(take_open, open_n, 0))
        m = jnp.where(n_sel > 0, d2_sum / jnp.maximum(n_sel, 1), inf)

        # winding number (piecewise constant, no gradient)
        ac, bc = jax.lax.stop_gradient(ac), jax.lax.stop_gradient(bc)
        x1, y1 = ac[:, 0:1], ac[:, 1:2]
        x2, y2 = bc[:, 0:1], bc[:, 1:2]
        is_left = (x2 - x1) * (py - y1) - (px - x1) * (y2 - y1)
        up = (y1 <= py) & (y2 > py) & (is_left > 0)
        down = (y2 <= py) & (y1 > py) & (is_left < 0)
        w = seg_sum(up.astype(jnp.int32) - down.astype(jnp.int32), wid)
        w = w.at[0].add(jnp.where(cont_in, open_w, 0))

        # fold every polygon that ends in this chunk into the union
        last = n_local - 1
        closed = (k < n_local) & ~((k == last) & cont_out)
        dist = jnp.sqrt(m + 1e-12)
        sd_vertex = jnp.where(w != 0, -dist, dist)
        on_edge = (n_int == n_sel) & (n_sel > 0)
        lin_k = lin_sum / jnp.maximum(n_int, 1)
        # the edge line's side is the polygon's side only near its boundary: with
        # seam edges excluded, the nearest remaining edge can face away from a
        # point beyond the seam, so away from the line the winding sets the sign
        lin_val = jax.lax.stop_gradient(lin_k)
        agrees = ((lin_val < 0) == (w != 0)) | (jnp.abs(lin_val) <= sign_tol)
        sd_k = jnp.where(on_edge & agrees, lin_k, sd_vertex)
        sd = jnp.minimum(sd, jnp.min(jnp.where(closed, sd_k, inf), axis=0))
        keep = cont_out
        return (
            sd,
            jnp.where(keep, m[last], inf),
            jnp.where(keep, lin_sum[last], 0.0),
            jnp.where(keep, n_int[last], 0),
            jnp.where(keep, n_sel[last], 0),
            jnp.where(keep, w[last], 0),
        ), None

    n_points = points.shape[0]
    zeros_i = jnp.zeros((n_points,), dtype=jnp.int32)
    init = (
        jnp.full((n_points,), jnp.inf, dtype=points.dtype),
        jnp.full((n_points,), jnp.inf, dtype=points.dtype),
        jnp.zeros((n_points,), dtype=points.dtype),
        zeros_i,
        zeros_i,
        zeros_i,
    )
    xs = (
        a,
        b,
        orient.astype(points.dtype),
        jnp.asarray(lay["wind_lid"]),
        jnp.asarray(lay["dist_lid"]),
        jnp.asarray(lay["n_local"]),
        jnp.asarray(lay["cont_in"]),
        jnp.asarray(lay["cont_out"]),
    )
    (sd, *_), _ = jax.lax.scan(body, init, xs)
    return sd

def pack_polygons(polygons: Sequence[Any]) -> tuple[Array, tuple[int, ...]]:
    """Concatenates (N_i, 2) vertex arrays into one (sum N_i, 2) array.

    Returns:
        vertices: the packed (differentiable) vertices.
        sizes: number of vertices per polygon (the static topology).
    """
    polys = [jnp.asarray(p) for p in polygons]
    if not polys:
        raise ValueError("Need at least one polygon")
    for p in polys:
        if p.ndim != 2 or p.shape[1] != 2 or p.shape[0] < 3:
            raise ValueError(f"Polygons must be (N >= 3, 2) arrays, got {p.shape}")
    return jnp.concatenate(polys, axis=0), tuple(int(p.shape[0]) for p in polys)

def split_t_junctions(
    vertices: np.ndarray, sizes: Sequence[int], tol: float = 1e-6
) -> tuple[np.ndarray, tuple[int, ...]]:
    """Inserts into every edge the polygon vertices that lie on it (T-junctions).

    Where a polygon corner lies inside another edge, e.g. a taper end in the
    middle of the side of an MMI, the two polygons share only part of that edge,
    so the shared part is not recognized as a seam. Splitting the long edge at
    the corner makes the shared part an exactly coincident (reversed) edge, which
    :func:`_edge_topology` then excludes from the distance. Inserted points lie
    on their edges, so the filled region does not change.

    Static (numpy, from concrete vertices); the inserted points copy vertices, so
    with ``vertices[index]`` they follow those vertices when they move, which is
    right as long as they stay inside their edges (see :func:`topology_changes`).
    A vertex up to ``tol`` off an edge is inserted too: the edge then bends
    through it by at most ``tol``, where leaving it out would leave a slit
    between the polygons (layout snapping moves a vertex off a slanted edge it
    was on).

    Args:
        vertices: (E, 2) packed polygon vertices.
        sizes: vertices per polygon.
        tol: distance (same units as ``vertices``) within which a vertex lies on
            an edge, and within which it is at the edge's end (then the two
            polygons share that vertex instead).

    Returns:
        index: indices into ``vertices`` of the split polygons' vertices.
        sizes: vertices per split polygon.
    """
    return _apply_inserts(_t_junction_inserts(vertices, sizes, tol), sizes)

def gaussian_random_field(
    shape: tuple[int, int], correlation_length: float | tuple[float, float], key: Array
) -> Array:
    """Unit-variance Gaussian random field with covariance exp(-r^2 / L^2).

    White noise filtered in Fourier space on a grid padded by 3 L on each side,
    so the field is not periodic over ``shape``. Lengths are in cells; a pair
    (Lx, Ly) is one length on cells of different sizes along x and y
    (covariance ``exp(-(x / Lx)^2 - (y / Ly)^2)`` in cells).
    """
    if isinstance(correlation_length, tuple):
        lx, ly = correlation_length
        pad = (math.ceil(3 * lx), math.ceil(3 * ly))
        n = (shape[0] + 2 * pad[0], shape[1] + 2 * pad[1])
        w = jax.random.normal(key, n)
        kx = 2 * jnp.pi * jnp.fft.fftfreq(n[0])
        ky = 2 * jnp.pi * jnp.fft.fftfreq(n[1])
        h = jnp.exp(-(kx[:, None] ** 2 * lx**2 + ky[None, :] ** 2 * ly**2) / 8)
        f = jnp.fft.ifft2(jnp.fft.fft2(w) * h).real / jnp.sqrt(jnp.mean(h**2))
        return f[pad[0] : pad[0] + shape[0], pad[1] : pad[1] + shape[1]]
    pad = math.ceil(3 * correlation_length)
    n = (shape[0] + 2 * pad, shape[1] + 2 * pad)
    w = jax.random.normal(key, n)
    kx = 2 * jnp.pi * jnp.fft.fftfreq(n[0])
    ky = 2 * jnp.pi * jnp.fft.fftfreq(n[1])
    k2 = kx[:, None] ** 2 + ky[None, :] ** 2
    h = jnp.exp(-k2 * correlation_length**2 / 8)  # sqrt of the power spectrum
    f = jnp.fft.ifft2(jnp.fft.fft2(w) * h).real / jnp.sqrt(jnp.mean(h**2))
    return f[pad : pad + shape[0], pad : pad + shape[1]]

def align_material_poles(materials: dict[str, fdtdx.Material]) -> dict[str, fdtdx.Material]:
    """Gives every material the same ordered pole slots.

    The slots are the distinct (omega_0, gamma) pairs of all poles of all
    materials, in order of first appearance; poles of one material that share
    (omega_0, gamma) are merged (their K add, which is exact). A material gets
    its own K in its slots and 0 elsewhere, so the discrete coefficients c1, c2
    of a slot are the same for every material and only c3 (proportional to K)
    differs. Mixing materials in a voxel then reduces to mixing c3, which mixes
    the susceptibilities exactly: ``chi = sum_m w_m chi_m``.

    Materials without poles stay unchanged when no material is dispersive.
    Permittivity (eps_inf), permeability and conductivities are kept, so
    fdtdx's material order does not change. Idempotent.
    """
    slots: list[tuple[float, float]] = []
    for m in materials.values():
        if m.dispersion is None:
            continue
        for pole in m.dispersion.poles:
            key = (float(pole.omega_0), float(pole.gamma))
            if key not in slots:
                slots.append(key)
    if not slots:
        return dict(materials)
    out = {}
    for name, m in materials.items():
        strength = [0.0] * len(slots)
        if m.dispersion is not None:
            for pole in m.dispersion.poles:
                strength[slots.index((float(pole.omega_0), float(pole.gamma)))] += float(
                    pole.coupling_sq
                )
        poles = tuple(
            SlotPole(resonance=w0, damping_rate=g, strength=k)
            for (w0, g), k in zip(slots, strength, strict=True)
        )
        out[name] = fdtdx.Material(
            permittivity=m.permittivity,
            permeability=m.permeability,
            electric_conductivity=m.electric_conductivity,
            magnetic_conductivity=m.magnetic_conductivity,
            dispersion=DispersionModel(poles=poles),
        )
    return out

def warn_reversible_losses(
    materials: dict[str, fdtdx.Material], config: SimulationConfig, name: str, conductivity: bool
) -> None:
    """Warns about lossy materials when gradients use fdtdx's reversible mode.

    The reversible backward pass runs the update backwards in time, which
    amplifies round-off in lossy voxels: by ~exp(gamma T / 2) for a damped
    Lorentz pole, ~exp(gamma T) for a Drude pole, and per step for a
    conductivity. Use ``GradientConfig(method="checkpointed")`` for those.
    """
    gradient_config = config.gradient_config
    if gradient_config is None or gradient_config.method != "reversible":
        return
    lossy = []
    for k, m in materials.items():
        if conductivity and m.is_electrically_conductive:
            lossy.append(f"{k} (conductivity)")
        if m.dispersion is not None:
            exponent = max(
                (
                    p.gamma * config.time * (1.0 if p.omega_0 == 0 else 0.5)
                    for p in m.dispersion.poles
                    if p.coupling_sq != 0
                ),
                default=0.0,
            )
            if exponent > 9.0:  # round-off amplified more than ~1e4 times
                lossy.append(f"{k} (pole damping: error growth ~exp({exponent:.0f}))")
    if lossy:
        warnings.warn(
            f"{name!r} contains lossy materials {lossy}; fdtdx's reversible gradients "
            "reconstruct the fields backwards in time, which amplifies round-off there. "
            'Use GradientConfig(method="checkpointed") for accurate gradients.',
            stacklevel=3,
        )