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_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)

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 get_basic_mass_coordinates_from_spatial_coordinates(
    self, basic_coordinates: npt.NDArray
) -> npt.NDArray:
    """Computes mass coordinates matching SPIDER's definition.

    SPIDER's mass coordinate (eos_adamswilliamson.c:296-311):

        xi(r)^3 = r_core^3 + 3 * M_AW(r_core, r) / rho_avg_mantle

    where M_AW is the A-W mass integral from r_core to r (without
    4pi), and rho_avg_mantle is the mantle-only average density.
    At the CMB: xi = r_core. At the surface: xi = r_surface
    (by construction of rho_avg_mantle).

    Args:
        Basic spatial coordinates

    Returns:
        Basic mass coordinates
    """
    r_core = basic_coordinates[0, 0]

    # xi^3 at CMB = r_core^3
    basic_mass_coordinates = np.zeros_like(basic_coordinates)
    basic_mass_coordinates[:, :] = np.power(r_core, 3.0)

    # Cumulative mantle mass contribution.
    # staggered_effective_density * (r^3_outer - r^3_inner) gives
    # mass_shell / (4/3*pi). Dividing by rho_avg (which is
    # M_no4pi / V_no4pi = M_no4pi / ((r_s^3-r_c^3)/3)) gives
    # the correct xi^3 increment: 3 * M_shell_no4pi / rho_avg.
    basic_volumes = np.power(basic_coordinates[1:, 0], 3.0) - np.power(
        basic_coordinates[:-1, 0], 3.0
    )
    for i in range(1, self.settings.number_of_nodes):
        basic_mass_coordinates[i:, :] += (
            self.staggered_effective_density[i - 1, :]
            * basic_volumes[i - 1]
            / self._planet_density
        )

    return np.power(basic_mass_coordinates, 1.0 / 3.0)

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 get_staggered_spatial_coordinates_from_mass_coordinates(
    self, staggered_mass_coordinates: npt.NDArray
) -> npt.NDArray:
    """Computes the staggered spatial coordinates from staggered mass coordinates.

    Args:
        Staggered mass coordinates

    Returns:
        Staggered spatial coordinates
    """

    # Initialise the staggered spatial coordinate to the inner boundary
    staggered_coordinates = np.ones_like(staggered_mass_coordinates) * np.power(
        self.settings.inner_radius, 3.0
    )

    # Add first half cell contribution
    staggered_coordinates += (
        self._planet_density
        * (
            np.power(staggered_mass_coordinates[0, :], 3.0)
            - np.power(self.basic.mass_radii[0, :], 3.0)
        )
        / self.staggered_effective_density[0, :]
    )

    # Get spatial coordinates by adding individual cell contributions to the mantle mass
    shell_effective_density = 0.5 * (
        self.staggered_effective_density[1:, :] + self.staggered_effective_density[:-1, :]
    )
    shell_mass_volumes = np.power(staggered_mass_coordinates[1:, :], 3.0) - np.power(
        staggered_mass_coordinates[:-1, :], 3.0
    )
    for i in range(1, self.settings.number_of_nodes - 1):
        staggered_coordinates[i:, :] += (
            self._planet_density
            * shell_mass_volumes[i - 1, :]
            / shell_effective_density[i - 1, :]
        )

    return np.power(staggered_coordinates, 1.0 / 3.0)

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 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)