Source code for magpylib_jax.magnet.cylinder

"""Differentiable cylinder magnet source."""

from __future__ import annotations

import jax.numpy as jnp

from magpylib_jax._types import ArrayLike
from magpylib_jax.constants import MU0
from magpylib_jax.core.base import BaseSource, MagpylibMissingInput




class Cylinder(BaseSource):
    """Homogeneously polarized cylinder with diameter-height dimensions."""

    _source_type = "cylinder"

    def __init__(
        self,
        dimension: ArrayLike | None = None,
        polarization: ArrayLike | None = None,
        magnetization: ArrayLike | None = None,
        position: ArrayLike = (0.0, 0.0, 0.0),
        orientation: ArrayLike | None = None,
        style=None,
        style_label: str | None = None,
        **kwargs,
    ) -> None:
        self.dimension = dimension
        self.polarization = polarization
        self.magnetization = magnetization
        super().__init__(
            position=position,
            orientation=orientation,
            style=style,
            style_label=style_label,
            **kwargs,
        )

    @property
    def _polarization(self) -> jnp.ndarray:
        if self.polarization is not None:
            return jnp.asarray(self.polarization, dtype=float)
        if self.magnetization is not None:
            return MU0 * jnp.asarray(self.magnetization, dtype=float)
        raise MagpylibMissingInput("Input polarization of Cylinder must be set.")

    @property
    def volume(self) -> float:
        if self.dimension is None:
            return 0.0
        diameter, height = jnp.asarray(self.dimension, dtype=float)
        radius = diameter / 2.0
        return float(jnp.pi * radius * radius * height)

    @property
    def centroid(self) -> jnp.ndarray:
        return jnp.asarray(self.position, dtype=float)

    def _require_inputs(self) -> None:
        if self.dimension is None:
            raise MagpylibMissingInput("Input dimension of Cylinder must be set.")
        if self.polarization is None and self.magnetization is None:
            raise MagpylibMissingInput("Input polarization of Cylinder must be set.")

    def _field_kwargs(self) -> dict:
        """Geometry + excitation arguments for the field engine."""
        return {"dimension": self.dimension, "polarization": self._polarization}