Mesh

def get_density(self, pressure: FloatOrArray) -> npt.NDArray:
    r"""Computes density from pressure (SPIDER parity):

    $$
    \rho(P) = \rho_s + P \beta / g
    $$

    Matches SPIDER's eos_adamswilliamson.c:GetRho (line 116):
    ``rho = rhos - P * beta / gravity`` where gravity is negative.

    Args:
        pressure: Pressure

    Returns:
        Density
    """
    density: npt.NDArray = (
        self._surface_density + pressure * self._beta / self._gravitational_acceleration
    )

    return density

def get_density_from_radii(self, radii: FloatOrArray) -> FloatOrArray:
    r"""Computes density from radii using SPIDER-parity exponential form:

    $$
        \rho(r) = \rho_s \exp(\beta (R_s - r))
    $$

    where $\rho_s$ is surface density, $\beta$ is the Adams-Williamson
    parameter [1/m], and $R_s$ is the surface radius. Matches SPIDER's
    eos_adamswilliamson.c via the chain rho(P(r)).

    Args:
        radii: Radii

    Returns
        Density
    """
    depth = self._outer_boundary - radii
    density: FloatOrArray = self._surface_density * np.exp(self._beta * depth)

    return density

def get_effective_density(self, radii) -> npt.NDArray:
    r"""
    Computes effective density using density rho(r) integration
    over a spherical shell bounded by radii

    Args:
        radii: Radii array

    Returns:
        Effective Density array
    """

    mass_shell = self.get_mass_within_shell(radii)
    volume_shell = 4 / 3 * np.pi * (np.power(radii[1:], 3.0) - np.power(radii[:-1], 3.0))

    return mass_shell / volume_shell

def get_mass_element(self, radii: FloatOrArray) -> npt.NDArray:
    r"""Computes the mass element:

    $$
    \frac{\delta m}{\delta r} = 4 \pi r^2 \rho
    $$

    where $\delta m$ is the mass element, $r$ is radius, and $\rho$ is
    density.

    Args:
        radii: Radii

    Returns:
        The mass element at radii
    """
    mass_element: npt.NDArray = (
        4 * np.pi * np.square(radii) * self.get_density_from_radii(radii)
    )

    return mass_element

def get_mass_within_radii(self, radii: FloatOrArray) -> npt.NDArray:
    r"""Computes mass within radii using SPIDER-parity exponential form:

    $$
    m(r) = 4\pi \int r^2 \rho_s \exp(\beta (R_s - r)) dr
    $$

    The antiderivative is (matching SPIDER eos_adamswilliamson.c:191):

    $$
    M(r) = 4\pi \rho(r) \left( -\frac{2}{\beta^3} - \frac{r^2}{\beta}
           - \frac{2r}{\beta^2} \right)
    $$

    where rho(r) = rhos * exp(beta*(R_s - r)).

    Args:
        radii: Radii

    Returns:
        Mass within radii (from inner_boundary to radii)
    """
    beta = self._beta

    def mass_integral(r: FloatOrArray) -> npt.NDArray:
        rho = self.get_density_from_radii(r)
        return 4.0 * np.pi * rho * (-2.0 / beta**3 - r**2 / beta - 2.0 * r / beta**2)

    mass: npt.NDArray = mass_integral(radii) - mass_integral(self._inner_boundary)

    return mass

def get_mass_within_shell(self, radii: npt.NDArray) -> npt.NDArray:
    """Computes the mass within spherical shells bounded by radii.

    Args:
        radii: Radii

    Returns:
        Mass within the bounded spherical shells
    """
    mass: npt.NDArray = self.get_mass_within_radii(radii[1:]) - self.get_mass_within_radii(
        radii[:-1]
    )

    return mass

def get_pressure_from_radii(self, radii: FloatOrArray) -> npt.NDArray:
    r"""Computes pressure from radii using SPIDER-parity exponential form:

    $$
    P(r) = \frac{\rho_s g}{\beta} \left( \exp(\beta (R_s - r)) - 1 \right)
          + P_{surf}
    $$

    Matches SPIDER's eos_adamswilliamson.c:GetPressureFromRadius.
    Note: g is positive here (unlike SPIDER where gravity is negative).

    Parameters
    ----------
    radii : FloatOrArray
        Radii at which to compute pressure.

    Returns
    -------
    npt.NDArray
        Pressure at the given radii.
    """
    depth = self._outer_boundary - radii
    pressure: npt.NDArray = (
        self._surface_density
        * self._gravitational_acceleration
        / self._beta
        * (np.exp(self._beta * depth) - 1.0)
        + self._surface_pressure
    )

    return pressure

def get_pressure_gradient(self, pressure: FloatOrArray) -> npt.NDArray:
    r"""Computes the pressure gradient:

    $$
    \frac{dP}{dr} = -g \rho
    $$

    where $\rho$ is density, $P$ is pressure, and  $g$ is gravitational
    acceleration.

    Args:
        pressure: Pressure

    Returns:
        Pressure gradient
    """
    dPdr: npt.NDArray = -self._gravitational_acceleration * self.get_density(pressure)

    return dPdr

def get_radii_from_pressure(self, pressure: FloatOrArray) -> npt.NDArray:
    r"""Computes radii from pressure (SPIDER parity):

    $$
    r(P) = R_s - \frac{1}{\beta} \ln\left(1 + \frac{P \beta}{\rho_s g}\right)
    $$

    Inverse of get_pressure_from_radii.

    Args:
        pressure: Pressure

    Returns:
        Radii
    """
    radii: npt.NDArray = (
        self._outer_boundary
        - np.log(
            1.0
            + (pressure - self._surface_pressure)
            * self._beta
            / (self._surface_density * self._gravitational_acceleration)
        )
        / self._beta
    )

    return radii

def set_basic_density(
    self,
    basic_radii: npt.NDArray,
) -> None:
    """Set basic-node density based on basic radii.

    Refreshing basic_density from the post-mass-coordinate radii keeps
    get_dxidr_basic consistent with the solved mesh.
    """
    self._basic_density = self.get_density_from_radii(basic_radii)

def set_basic_pressure(
    self,
    basic_radii: npt.NDArray,
) -> None:
    """Set basic pressure based on basic radii."""
    self._basic_pressure = self.get_pressure_from_radii(basic_radii)

def set_staggered_pressure(
    self,
    staggered_radii: npt.NDArray,
) -> None:
    """Set staggered pressure based on staggered radii."""
    self._staggered_pressure = self.get_pressure_from_radii(staggered_radii)

@abstractmethod
def get_mass_within_radii(self, radii: FloatOrArray) -> npt.NDArray:
    """Cumulative mass [kg] inside ``radii`` from a class-specific zero point.

    The result has the same shape as ``radii``.
    """

def set_staggered_effective_density(self, basic_radii: npt.NDArray) -> None:
    """Set each cell's effective density to its mass over its volume.

    Parameters
    ----------
    basic_radii : npt.NDArray
        Basic-node radii [m], shape (n, 1), the faces of the cells.
    """
    mass = np.diff(self.get_mass_within_radii(basic_radii), axis=0)
    volume = 4.0 / 3.0 * np.pi * np.diff(np.power(basic_radii, 3.0), axis=0)
    self._staggered_effective_density = mass / volume

def d_dr_at_basic_nodes(self, staggered_quantity: npt.NDArray) -> npt.NDArray:
    """Determines d/dr at the basic nodes of a quantity defined at the staggered nodes.

    Args:
        staggered_quantity: A quantity defined at the staggered nodes.

    Returns:
        d/dr at the basic nodes
    """
    d_dr_at_basic_nodes: npt.NDArray = self._d_dr_transform.dot(staggered_quantity)
    logger.debug('d_dr_at_basic_nodes = %s', d_dr_at_basic_nodes)

    return d_dr_at_basic_nodes

def entropy_at_basic_nodes(self, S: npt.NDArray) -> npt.NDArray:
    """Entropy at the basic nodes, with the top node from ``top_node_entropy`` when limited.

    Raises
    ------
    RuntimeError
        When a limited node with finite cells is not between the top cell and the linear
        node, or leaves the table with the top cell in it and is not the linear node within
        4 ulp of the table.
    """
    S = np.ravel(S)
    S_basic = self.quantity_at_basic_nodes(S)
    if self.top_node == 'limited':
        lo, hi = self.top_bounds
        node = float(top_node_entropy(S, self.top_ratio, lo, hi, linear=S_basic[-1]))
        ends = sorted((S[-1], S_basic[-1]))
        tol = 1e-12 * max(abs(ends[0]), abs(ends[1]), 1.0)
        # the linear row may lie outside by rounding: 4 ulp, 2 measured at ratio < 2
        near = min(node - lo, hi - node) >= -4 * np.spacing(abs(node))
        if np.isfinite(S_basic[-1]) and not (
            ends[0] - tol <= node <= ends[1] + tol
            and (
                not lo <= S[-1] <= hi or lo <= node <= hi or (node == S_basic[-1] and near)
            )
        ):
            raise RuntimeError(
                f'limited top node {node!r} is outside [{ends[0]!r}, {ends[1]!r}] (top cell, '
                f'linear node) or the table [{lo!r}, {hi!r}] that holds the top cell'
            )
        S_basic[-1] = node
    return S_basic

def get_constant_spacing(self) -> npt.NDArray:
    """Constant radius spacing across the mantle

    Returns:
        Radii with constant spacing as a column vector
    """
    radii: npt.NDArray = np.linspace(
        self.settings.inner_radius,
        self.settings.outer_radius,
        self.settings.number_of_nodes,
    )
    radii = np.atleast_2d(radii).T

    return radii

def get_dxidr_basic(self) -> npt.NDArray:
    """Computes dxidr at basic nodes."""

    dxidr = (
        self.eos.basic_density
        / self._planet_density
        * np.power(self.basic.radii, 2.0)
        / np.power(self.basic.mass_radii, 2.0)
    )

    return dxidr

def get_planet_density(self, basic_coordinates: npt.NDArray) -> float:
    """Computes the mantle average density for mass-coordinate mapping.

    Matches SPIDER's ``EOSAdamsWilliamson_GetMassCoordinateAverageRho``
    (eos_adamswilliamson.c:249-267): the average density is computed
    over the mantle shell only (r_core to r_surface), NOT including
    the core. This ensures the mass-coordinate xi grid maps to the
    same physical radii as SPIDER's Newton-solved mesh.

    The previous implementation used whole-planet density (core +
    mantle) / r_surface^3, which produced ~3% node position offsets
    and cascading 12% pressure / 20% density / 17% flux differences.

    Args:
        Basic spatial coordinates

    Returns:
        Mantle average density (for mass-coordinate normalization)
    """
    r_core = basic_coordinates[0, 0]
    r_surf = basic_coordinates[-1, 0]
    # Use the analytic A-W mass integral (without 4pi, SPIDER
    # convention) for the average density. The discrete sum
    # (rho * delta_r^3) is inconsistent because the 4pi/3
    # factors don't cancel between the volume elements and
    # the effective density definition.
    M_4pi = (
        self.eos.get_mass_within_radii(np.array([[r_surf]]))
        - self.eos.get_mass_within_radii(np.array([[r_core]]))
    ).item()
    M_no4pi = M_4pi / (4.0 * np.pi)
    mantle_volume_no4pi = (np.power(r_surf, 3.0) - np.power(r_core, 3.0)) / 3.0
    mantle_avg_density = M_no4pi / mantle_volume_no4pi
    return float(mantle_avg_density)

def quantity_at_basic_nodes(self, staggered_quantity: npt.NDArray) -> npt.NDArray:
    """Determines a quantity at the basic nodes that is defined at the staggered nodes.

    Uses backward and forward differences at the inner and outer radius, respectively, to
    obtain the quantity values of the basic nodes at the innermost and outermost nodes.
    When using temperature boundary conditions, values at outer boundaries will be overwritten.
    When using flux boundary conditions, values at outer boundaries will be used to provide
    estimate of individual components of heat fluxes though the total heat flux is imposed.

    Args:
        staggered_quantity: A quantity defined at the staggered nodes

    Returns:
        The quantity at the basic nodes
    """
    quantity_at_basic_nodes: npt.NDArray = self._quantity_transform.dot(staggered_quantity)
    logger.debug('quantity_at_basic_nodes = %s', quantity_at_basic_nodes)

    return quantity_at_basic_nodes

def quantity_at_staggered_nodes(self, basic_quantity: npt.NDArray) -> npt.NDArray:
    """Determines a quantity at the staggered nodes that is defined at the basic nodes.

    Staggered nodes are always located at cell centers, whatever the mesh.

    Args:
        basic_quantity: A quantity defined at the basic nodes

    Returns:
        The quantity at the staggered nodes
    """
    quantity_at_staggered_nodes: npt.NDArray = 0.5 * (
        basic_quantity[:-1, ...] + basic_quantity[1:, ...]
    )
    logger.debug('quantity_at_staggered_nodes = %s', quantity_at_staggered_nodes)

    return quantity_at_staggered_nodes

def get_mass_within_radii(self, radii: FloatOrArray) -> npt.NDArray:
    r"""4pi-included cumulative mass M(r), anchored at the inner EOS radius.

    Mesh.get_planet_density and the mass-coordinate brentq loop use only
    differences M(r2) - M(r1), so the anchor cancels.
    """
    return np.asarray(self._interp_mass(np.asarray(radii)))

def get_pressure_from_radii(self, radii: FloatOrArray) -> npt.NDArray:
    """Computes pressure from radii by interpolating the user-supplied EOS table.

    Args:
        radii: Radii at which to compute pressure.

    Returns:
        Pressure at the given radii.
    """
    pressure: npt.NDArray = self._interp_pressure(radii).reshape(-1, 1)

    return pressure

def set_basic_density(
    self,
    basic_radii: npt.NDArray,
) -> None:
    """Set basic-node density based on basic radii.

    Interpolate the user density table at the post-mass-coordinate basic
    radii so get_dxidr_basic reads density on the solved mesh.
    """
    self._basic_density = self._interp_density(basic_radii).reshape(-1, 1)

def set_basic_pressure(
    self,
    basic_radii: npt.NDArray,
) -> None:
    """Set basic pressure based on basic radii."""
    self._basic_pressure = self.get_pressure_from_radii(basic_radii)

def set_staggered_pressure(
    self,
    staggered_radii: npt.NDArray,
) -> None:
    """Set staggered pressure based on staggered radii."""
    self._staggered_pressure = self._interp_pressure(staggered_radii).reshape(-1, 1)